119 lines
2.8 KiB
Python
119 lines
2.8 KiB
Python
#!/usr/bin/env python
|
|
# coding: utf-8
|
|
|
|
# In[1]:
|
|
|
|
|
|
import pandas as pd
|
|
import csv
|
|
import regex as re
|
|
import kenlm
|
|
from english_words import english_words_alpha_set
|
|
from math import log10
|
|
from nltk import trigrams, word_tokenize
|
|
|
|
|
|
# In[2]:
|
|
|
|
|
|
default_pred = 'to:0.02 be:0.02 the:0.02 or:0.01 not:0.01 and:0.01 a:0.01 :0.9'
|
|
|
|
|
|
# In[3]:
|
|
|
|
|
|
def preprocess(text):
|
|
text = text.lower().replace('-\\n', '').replace('\\n', ' ')
|
|
return re.sub(r'\p{P}', '', text)
|
|
|
|
|
|
# In[4]:
|
|
|
|
|
|
train_data = pd.read_csv('train/in.tsv.xz', sep='\t', on_bad_lines='skip', header=None, quoting=csv.QUOTE_NONE,
|
|
nrows=20000)
|
|
train_labels = pd.read_csv('train/expected.tsv', sep='\t', on_bad_lines='skip', header=None,
|
|
quoting=csv.QUOTE_NONE, nrows=20000)
|
|
data = pd.concat([train_data, train_labels], axis=1)
|
|
data=train_data[6] + train_data[0] + train_data[7]
|
|
data = data.apply(preprocess)
|
|
|
|
with open("train_file.txt", "w+") as f:
|
|
for text in data:
|
|
f.write(text + "\n")
|
|
|
|
|
|
# In[5]:
|
|
|
|
|
|
KENLM_BUILD_PATH='../kenlm/kenlm/build'
|
|
get_ipython().system('$KENLM_BUILD_PATH/bin/lmplz -o 4 < train_file.txt > model.arpa')
|
|
get_ipython().system('rm train_file.txt')
|
|
|
|
|
|
# In[6]:
|
|
|
|
|
|
model = kenlm.Model("model.arpa")
|
|
|
|
|
|
# In[7]:
|
|
|
|
|
|
def predict(before, after):
|
|
best_scores = []
|
|
for word in english_words_alpha_set:
|
|
text = ' '.join([before, word, after])
|
|
text_score = model.score(text, bos=False, eos=False)
|
|
if len(best_scores) < 12:
|
|
best_scores.append((word, text_score))
|
|
else:
|
|
is_better = False
|
|
worst_score = None
|
|
for score in best_scores:
|
|
if not worst_score:
|
|
worst_score = score
|
|
else:
|
|
if worst_score[1] > score[1]:
|
|
worst_score = score
|
|
if worst_score[1] < text_score:
|
|
best_scores.remove(worst_score)
|
|
best_scores.append((word, text_score))
|
|
probs = sorted(best_scores, 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
|
|
|
|
|
|
# In[8]:
|
|
|
|
|
|
def make_prediction(path, result_path):
|
|
data = pd.read_csv(path, sep='\t', on_bad_lines='skip', header=None, quoting=csv.QUOTE_NONE)
|
|
with open(result_path, 'w', encoding='utf-8') as file_out:
|
|
for _, row in data.iterrows():
|
|
before, after = word_tokenize(preprocess(str(row[6]))), word_tokenize(preprocess(str(row[7])))
|
|
if len(before) < 2 or len(after) < 2:
|
|
prediction = default_pred
|
|
else:
|
|
prediction = predict(before[-1], after[0])
|
|
file_out.write(prediction + '\n')
|
|
|
|
|
|
# In[9]:
|
|
|
|
|
|
make_prediction("dev-0/in.tsv.xz", "dev-0/out.tsv")
|
|
|
|
|
|
# In[10]:
|
|
|
|
|
|
make_prediction("test-A/in.tsv.xz", "test-A/out.tsv")
|
|
|
|
|
|
|
|
|