challenging-america-word-ga.../run.py
2022-04-24 22:05:42 +02:00

126 lines
3.4 KiB
Python

from cmath import log10
import csv
import pandas as pd
import regex as re
import os
import kenlm
from nltk import word_tokenize
from collections import Counter, defaultdict
from english_words import english_words_set
# nltk.download("punkt")
# train set
train_data = pd.read_csv(
"train/in.tsv.xz",
sep="\t",
error_bad_lines=False,
warn_bad_lines=False,
header=None,
quoting=csv.QUOTE_NONE,
nrows=100_000
)
# training labels
train_labels = pd.read_csv(
"train/expected.tsv",
sep="\t",
error_bad_lines=False,
warn_bad_lines=False,
header=None,
quoting=csv.QUOTE_NONE,
nrows=100_000
)
# dev set
dev_data = pd.read_csv(
"dev-0/in.tsv.xz",
sep="\t",
error_bad_lines=False,
warn_bad_lines=False,
header=None,
quoting=csv.QUOTE_NONE,
)
# test set
test_data = pd.read_csv(
"test-A/in.tsv.xz",
sep="\t",
error_bad_lines=False,
warn_bad_lines=False,
header=None,
quoting=csv.QUOTE_NONE,
)
def prepare_text(text):
text = text.lower().replace("-\\n", "").replace("\\n", " ")
text = re.sub(r"\p{P}", "", text)
return text
def predict(word1, word2):
predictions = []
for word in english_words_set:
sentence = word1 + ' ' + word + ' ' + word2
text_score = model.score(sentence, bos=False, eos=False)
if len(predictions) < 12:
predictions.append((word, text_score))
else:
worst_score = None
for score in predictions:
if not worst_score:
worst_score = score
else:
if worst_score[1] > score[1]:
worst_score = score
if worst_score[1] < text_score:
predictions.remove(worst_score)
predictions.append((word, text_score))
probs = sorted(predictions, key=lambda tup: tup[1], reverse=True)
pred_str = ''
for word, prob in probs:
pred_str += f'{word}:{prob} '
pred_str += f':{log10(0.99)}'
return pred_str
def write_output():
with open("dev-0/out.tsv", "w") as file:
for _, row in dev_data.iterrows():
text = prepare_text(str(row[7]))
words = word_tokenize(text)
if len(words) < 3:
prediction = "the:0.2 be:0.2 to:0.2 of:0.1 and:0.1 a:0.1 :0.1"
else:
prediction = predict(words[0], words[1])
file.write(prediction + "\n")
with open("test-A/out.tsv", "w") as file:
for _, row in test_data.iterrows():
text = prepare_text(str(row[7]))
words = word_tokenize(text)
if len(words) < 3:
prediction = "the:0.2 be:0.2 to:0.2 of:0.1 and:0.1 a:0.1 :0.1"
else:
prediction = predict(words[0], words[1])
file.write(prediction + "\n")
if __name__ == "__main__":
print("Preparing data...")
train_data = train_data[[6, 7]]
train_data = pd.concat([train_data, train_labels], axis=1)
train_data["final"] = train_data[6] + train_data[0] + train_data[7]
train = train_data[['final']]
with open("./train_data.txt", 'a') as f:
for _, row in train_data.iterrows():
text = prepare_text(str(row["final"]))
f.write(text + '\n')
print("Preparing model...")
os.system('sh ./kenlm.sh')
model=kenlm.Model("kenlm_model.binary")
print("Writing outputs...")
write_output()