#!/usr/bin/env python
# coding: utf-8

# In[1]:


import torch
from transformers import GPT2Tokenizer, GPT2LMHeadModel

device = torch.device('cuda')
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model: GPT2LMHeadModel = GPT2LMHeadModel.from_pretrained('gpt2', pad_token_id=tokenizer.eos_token_id)
model.to(device)

# In[2]:


import lzma


def read_xz_file(fname):
    with lzma.open(fname, mode='rt', encoding='utf-8') as f:
        return [line.strip() for line in f.readlines()]


# In[3]:


dev_input_raw = read_xz_file('dev-0/in.tsv.xz')
test_input_raw = read_xz_file('test-A/in.tsv.xz')


# In[4]:


def get_contexts(input_text):
    all_fields = input_text.replace(r'\n', ' ').split('\t')
    return {'left': all_fields[6], 'right': all_fields[7]}


dev_input_contexts = [get_contexts(input_text) for input_text in dev_input_raw]

# In[5]:


test_input_contexts = [get_contexts(input_text) for input_text in test_input_raw]

# In[6]:


from tqdm import tqdm

tokenizer.truncation_side = 'left'


def predict_words(dataset):
    preds = []
    for entry in tqdm(dataset):
        text = f"{entry['left']}"
        src = tokenizer.encode(text, return_tensors="pt", truncation=True).to(device)
        output = model.generate(src, max_length=len(src[0]) + 1, do_sample=True, top_k=0, temperature=0.8,
                                num_return_sequences=1, no_repeat_ngram_size=2)
        generated_word = tokenizer.decode(output[0], skip_special_tokens=True).split(' ')[-1]
        preds.append(f'{generated_word.strip()}:0.99 :0.01')
    return preds


# In[7]:


dev_preds = predict_words(dev_input_contexts)
with open('dev-0/out.tsv', 'w') as f:
    f.writelines(line + '\n' for line in dev_preds)

# In[8]:


test_preds = predict_words(test_input_contexts)
with open('test-A/out.tsv', 'w') as f:
    f.writelines(line + '\n' for line in test_preds)