challenging-america-word-gap-prediction/model.py

import lzma
from nltk.tokenize import word_tokenize
from nltk import trigrams
import string
from collections import defaultdict, Counter


trigrams_list = []
model = defaultdict(lambda: defaultdict(lambda: 0))


def preprocess(text):
    _text = text.lower().replace('\\n', ' ').strip()
    for character in _text:
        if character not in string.ascii_lowercase + ' ':
            _text = _text.replace(character, '')
    words = word_tokenize(_text)
    if len(words):
        return words
    return ['']


def predict(word_before, word_after):
    prob_list = model[(word_before, word_after)].items()
    predictions = []
    total = 0.0
    for key, value in prob_list:
        total += value
        predictions.append(f'{key}:{value}')
    if total == 0.0:
        return 'the:1.0'
    return ' '.join(predictions)


with open('train/in.tsv', 'w', encoding='utf-8') as file:
    print('dekompresja pliku')
    text = lzma.open('train/in.tsv.xz').read().decode('utf-8')
    file.write(text)

with open('train/in.tsv', encoding='utf-8') as file_in, open('train/expected.tsv', encoding='utf-8') as file_expected:
    for index, (line_in, expected) in enumerate(zip(file_in, file_expected)):
        if index % 1000 == 0:
            print('zbieranie trigramów', index)
        _, _, _, _, _, _, before, after = line_in.split('\t')
        before, expected, after = preprocess(before), preprocess(expected), preprocess(after)
        words = before + expected + after
        trigrams_list += trigrams(words, pad_right=True, pad_left=True)
        length = len(trigrams_list)
        print('zbieranie trigramów:', length)
        if length > 1000000:
            break

for index, trigram in enumerate(trigrams_list):
    if index % 100000 == 0:
        print('uczenie modelu', index)
    if not trigram[0] or not trigram[1] or not trigram[2]:
        continue
    model[(trigram[0], trigram[2])][trigram[1]] += 1
    if index == 999999:
        break


for index, words_1_3 in enumerate(model):
    if index % 100000 == 0:
        print('normalizacja', index)
    count = sum(model[words_1_3].values())
    for word_2 in model[words_1_3]:
        model[words_1_3][word_2] /= float(count)

with open('test-A/in.tsv', encoding='utf-8') as file_in, open('test-A/out.tsv', 'w', encoding='utf-8') as file_out:
    print('zapisywanie test-A')
    for line_in in file_in:
        _, _, _, _, _, _, before, after = line_in.split('\t')
        word_before_in, word_after_in = preprocess(before)[-1], preprocess(after)[0]
        file_out.write(predict(word_before_in, word_after_in) + '\n')

with open('dev-0/in.tsv', encoding='utf-8') as file_in, open('dev-0/out.tsv', 'w', encoding='utf-8') as file_out:
    print('zapisywanie dev-0')
    for line_in in file_in:
        _, _, _, _, _, _, before, after = line_in.split('\t')
        word_before_in, word_after_in = preprocess(before)[-1], preprocess(after)[0]
        file_out.write(predict(word_before_in, word_after_in) + '\n')

print('koniec')