kenlm

.gitignore

@@ -8,4 +8,5 @@
 .token
 geval
 *in.tsv
-train_file.txt
+train_file.txt
+model.arpa

run.py

@@ -1,11 +1,14 @@
 #%%
 import pandas as pd
-import kenlm
 import csv
+import os
+import kenlm
+from nltk import word_tokenize
+from collections import defaultdict, Counter
 
 #%%
 def clean(text):
-    text = str(text).lower().strip().replace("’", "'").replace('\\n', " ").replace("'t", " not").replace("'s", " is").replace("'ll", " will").replace("'m", " am").replace("'ve", " have").replace(",", "").replace("-", "")
+    text = str(text).lower().strip().replace("’", "'").replace('\\n', " ").replace("'t", " not").replace("'s", " is").replace("'ll", " will").replace("'m", " am").replace("'ve", " have").replace(",", "").replace("-", "").replace(".", "").replace("'", "".replace("”", ""))
     return text
 
 train_in = pd.read_csv("train/in.tsv.xz", sep='\t', header=None, encoding="UTF-8", on_bad_lines="skip", quoting=csv.QUOTE_NONE, nrows=300000)[[6, 7]]
@@ -14,26 +17,49 @@ data = pd.concat([train_in, train_expected], axis=1)
 data = data[6] + data[0] + data[7]
 data = data.apply(clean)
 
-with open("train_file.txt", "w+") as f:
-    for text in data:
-        f.write(text + "\n")
-
-#%%
-KENLM_BUILD_PATH='../kenlm/build'
-!$KENLM_BUILD_PATH/bin/lmplz -o 4 < train_file.txt > model.arpa
-!rm train_file.txt
-
+if not os.path.isfile('train_file.txt'):
+    with open("train_file.txt", "w+") as f:
+        for text in data:
+            f.write(text + "\n")
 
 #%%
+#!../kenlm/build/bin/lmplz -o 4 < train_file.txt > model.arpa
 model = kenlm.Model("model.arpa")
+
+#%%
+import nltk
+nltk.download('punkt')
+
+#%%
+most_common = defaultdict(lambda: 0)
+for text in data:
+    words = word_tokenize(text)
+    if "d" in words:
+        words.remove("d")
+    for w in words:
+        most_common[w] += 1
+most_common = Counter(most_common).most_common(8000)
+
 #%%
 def predict(path, result_path):
     data = pd.read_csv(path, sep='\t', header=None, encoding="UTF-8", on_bad_lines="skip", quoting=csv.QUOTE_NONE)[7]
     with open(result_path, "w+", encoding="UTF-8") as f:
         for text in data:
-            #test
-            print(model.score(text, bos = True, eos = True))
-            break
+            result = {}
+            for word in most_common:
+                prob = model.score(text + f" {word[0]}")
+                result[word[0]] = prob
+            predictions = dict(Counter(result).most_common(6))
+            result = ""
+            for word, prob in predictions.items():
+                result += f"{word}:{prob} "
+            result = result.rstrip()
+            if len(result) == 0:
+                result = "a:0.2 the:0.2 to:0.2 of:0.1 and:0.1 of:0.1 :0.1"
+            f.write(result + "\n")
+            print(result)
+
 
 predict("dev-0/in.tsv.xz", "dev-0/out.tsv")
-predict("test-A/in.tsv.xz", "test-A/out.tsv")
+predict("test-A/in.tsv.xz", "test-A/out.tsv")
+# %%