transfix-mt/rapidfuzztest.ipynb

import copy
import nltk
import pandas as pd
import rapidfuzz
import time

from nltk.stem import WordNetLemmatizer
from rapidfuzz.fuzz import  partial_ratio
from rapidfuzz.utils import default_process


wl = WordNetLemmatizer()

glossary = pd.read_csv('mt-summit-corpora/glossary.tsv', sep='\t', header=None, names=['source', 'result'])

source_lemmatized = []
for word in glossary['source']:
    word = nltk.word_tokenize(word)
    source_lemmatized.append(' '.join([wl.lemmatize(x) for x in word]))

glossary['source_lem'] = source_lemmatized
glossary = glossary[['source', 'source_lem', 'result']]
glossary.set_index('source_lem')

# train_in_path = 'mt-summit-corpora/train/in.tsv'
# train_expected_path = 'mt-summit-corpora/train/expected.tsv'

train_in_path = 'mt-summit-corpora/dev-0/in.tsv'
train_expected_path = 'mt-summit-corpora/dev-0/expected.tsv'


file_pl = pd.read_csv(train_expected_path, sep='\t', header=None, names=['text'])

start_time = time.time_ns()
file_lemmatized = []
with open(train_in_path, 'r') as file:
    for line in file:
        if len(file_lemmatized) % 50000 == 0:
            print(len(file_lemmatized), end='\r')
        line = nltk.word_tokenize(line)
        file_lemmatized.append(' '.join([wl.lemmatize(x) for x in line]))

stop = time.time_ns()
timex = (stop - start_time) / 1000000000
print(timex)

0.194806501

THRESHOLD = 88

def is_injectable(sentence_pl, sequence):
    sen = sentence_pl.split()
    window_size = len(sequence.split())
    maxx = 0
    for i in range(len(sen) - window_size):
        current = rapidfuzz.fuzz.ratio(' '.join(sen[i:i + window_size]), sequence)
        if current > maxx:
            maxx = current
    if maxx >= THRESHOLD:
        return True
    else:
        return False

def get_injected(sentence, sequence, inject):
    sen = sentence.split()
    window_size = len(sequence.split())
    maxx = 0
    maxxi = 0
    for i in range(len(sen) - window_size + 1):
        current = rapidfuzz.fuzz.ratio(' '.join(sen[i:i + window_size]), sequence)
        if current >= maxx:
            maxx = current
            maxxi = i
    return ' '.join(sen[:maxxi + window_size]) + ' $' + inject + '$ ' + ' '.join(sen[maxxi + window_size:])

glossary['source_lem'] = [' ' + str(default_process(x)) + ' ' for x in glossary['source_lem']]

start_time = time.time_ns()
en = []
translation_line_counts = []
for line, line_pl in zip(file_lemmatized, file_pl['text'].values.tolist()):
    line = default_process(line)
    line_pl = default_process(line_pl)
    matchez = rapidfuzz.process.extract(query=line, choices=glossary['source_lem'], limit=5, score_cutoff=THRESHOLD, scorer=partial_ratio)
    if len(matchez) > 0:
        lines_added = 0
        for match in matchez:
            polish_translation = glossary.loc[lambda df: df['source_lem'] == match[0]]['result'].astype(str).values.flatten()[0]
            if is_injectable(line_pl, polish_translation):
                en.append(get_injected(line, match[0], polish_translation))
                lines_added += 1
        if lines_added == 0:
            en.append(line)
            lines_added = 1
        translation_line_counts.append(lines_added)
    else:
        translation_line_counts.append(1)
        en.append(line)


stop = time.time_ns()
timex = (stop - start_time) / 1000000000
print(timex)

6.904260614

translations = pd.read_csv(train_expected_path, sep='\t', header=0, names=['text'])
with open(train_expected_path + '.injected', 'w') as file_plx:
    for line, translation_line_ct in zip(translations['text'].values.tolist(), translation_line_counts):
        for i in range(translation_line_ct):
            file_plx.write(line + '\n')


with open(train_in_path + '.injected', 'w') as file_en:
    for e in en:
        file_en.write(e + '\n')