"""
spark-submit --executor-memory 5g --num-executors 10 --master yarn-client --files organic_label_counts.txt ./wikipedia/typograms/spark_predicted_label_counts.py "/data/wikipedia_plaintext" "/user/roman/predicted_label_counts"
"""
import sys
from pyspark import SparkContext
from kilogram.lang.tokenize import default_tokenize_func
from kilogram.dataset.edit_histories.wikipedia import line_filter

sc = SparkContext(appName="WikipediaPredictedLabelCounts")

wiki_plain = sc.textFile(sys.argv[1])

organic_label_dict = {}

for line in open("organic_label_counts.txt"):
    uri, label, counts = line.split('\t')
    if not label.islower():
        organic_label_dict[label.lower()] = uri
    organic_label_dict[label] = uri


def generate_ngrams(line):
    labels = []
    line = line.strip()
    for sentence in line_filter(' '.join(default_tokenize_func(line))):
        sentence = sentence.split()
        i = 0
        while i < len(sentence):
            for j in range(min(len(sentence), 20), i, -1):
                token = ' '.join(sentence[i:j])
                if i+1 == j and i == 0:
                    # if first word in sentence -> skip, could be wrong (Apple)
                    continue
                elif token in organic_label_dict:
                    labels.append(((token.lower(), organic_label_dict[token]), {token: 1}))
                    i = j-1
                    break
            i += 1
    return labels

def collect_tokens(value1, value2):
    for label, count in value1.items():
        if label in value2:
            value2[label] += count
        else:
            value2[label] = count
    return value2

wiki_predicted_labels = wiki_plain.flatMap(generate_ngrams).reduceByKey(collect_tokens)
wiki_predicted_labels.saveAsTextFile(sys.argv[2])

@@ -1727,16 +1727,443 @@
 tokens)%0A
+%0Adef printer(item):%0A    uri = item%5B0%5D%5B1%5D%0A    label_lower = item%5B0%5D%5B1%5D%0A    count_dict = item%5B1%5D%0A    count_lower = '0'%0A    count_normal = '0'%0A    if label_lower in count_dict:%0A        count_lower = count_dict%5Blabel_lower%5D%0A        del count_dict%5Blabel_lower%5D%0A    if len(count_dict) %3E 0:%0A        label_lower, count_normal = count_dict.items()%5B0%5D%0A    return uri + '%5Ct' + label_lower + '%5Ct' + str(count_normal)+','+str(count_lower)%0A%0A
 wiki_pre
@@ -2176,16 +2176,29 @@
 _labels.
+map(printer).
 saveAsTe

proper writer