challenging-america-word-gap-prediction/neural_network_simple_colab.ipynb

import itertools
import lzma
import regex as re
import torch
from torch import nn
from torch.utils.data import IterableDataset, DataLoader
from torchtext.vocab import build_vocab_from_iterator
import pickle
import os

from google.colab import drive
drive.mount('/content/drive')

Definitions

def clean_line(line: str):
    separated = line.split('\t')
    prefix = separated[6].replace(r'\n', ' ').strip()
    suffix = separated[7].replace(r'\n', ' ').strip()
    return prefix + ' ' + suffix

def get_words_from_line(line):
    line = clean_line(line)
    for word in line.split():
        yield word

def get_word_lines_from_file(file_name):
    with lzma.open(file_name, mode='rt', encoding='utf-8') as fid:
        for line in fid:
            yield get_words_from_line(line)


def look_ahead_iterator(gen):
    prev = None
    for item in gen:
        if prev is not None:
            yield (prev, item)
        prev = item

def predict(word: str, num_of_top: str) -> str:
    ixs = torch.tensor(vocab.forward([word])).to(device)
    out = model(ixs)
    top = torch.topk(out[0], num_of_top)
    top_indices = top.indices.tolist()
    top_probs = top.values.tolist()
    top_words = vocab.lookup_tokens(top_indices)
    zipped = list(zip(top_words, top_probs))
    if '<unk>' in [element[0] for element in zipped]:
        zipped = [(element[0] if element[0] != '<unk>' else '', element[1]) for element in zipped]
        zipped[-1] = ('', zipped[-1][1])
    else:
        zipped[-1] = ('', zipped[-1][1])
    return ' '.join([f'{element[0]}:{element[1]}' for element in zipped])

def execute(path):
    with lzma.open(f'{path}/in.tsv.xz', 'rt', encoding='utf-8') as f, \
        open(f'{path}/out.tsv', 'w', encoding='utf-8') as out:
        for line in f:
            prefix = line.split('\t')[6]
            left = prefix.replace(r'\n', ' ').split()[-1]
            result = predict(left, num_of_top)
            out.write(f"{result}\n")

class Bigrams(IterableDataset):
  def __init__(self, text_file, vocabulary_size):
      self.vocab = vocab
      self.vocab.set_default_index(self.vocab['<unk>'])
      self.vocabulary_size = vocabulary_size
      self.text_file = text_file

  def __iter__(self):
     return look_ahead_iterator(
         (self.vocab[t] for t in itertools.chain.from_iterable(get_word_lines_from_file(self.text_file))))

        
class SimpleBigramNeuralLanguageModel(nn.Module):
  def __init__(self, vocabulary_size, embedding_size):
      super(SimpleBigramNeuralLanguageModel, self).__init__()
      self.model = nn.Sequential(
          nn.Embedding(vocabulary_size, embedding_size),
          nn.Linear(embedding_size, vocabulary_size),
          nn.Softmax()
      )

  def forward(self, x):
      return self.model(x)

Parameters

vocab_size = 10000
embed_size = 250
batch_size = 5000
num_of_top = 10

Vocabulary building

if os.path.exists('./vocabulary.pickle'):
    with open('vocabulary.pickle', 'rb') as handle:
        vocab = pickle.load(handle)
else:
    vocab = build_vocab_from_iterator(
    get_word_lines_from_file('./drive/MyDrive/ColabNotebooks/america/train/in.tsv.xz'),
    max_tokens = vocab_size,
    specials = ['<unk>'])

    with open('vocabulary.pickle', 'wb') as handle:
        pickle.dump(vocab, handle, protocol=pickle.HIGHEST_PROTOCOL)

vocab.lookup_tokens([0, 1, 2, 3, 4, 4500])

Training

model = SimpleBigramNeuralLanguageModel(vocab_size, embed_size)
vocab.set_default_index(vocab['<unk>'])

#uczenie
from torch.utils.data import DataLoader

device = 'cuda'
train_dataset = Bigrams('./drive/MyDrive/ColabNotebooks/america/train/in.tsv.xz', vocab_size)
model = SimpleBigramNeuralLanguageModel(vocab_size, embed_size).to(device)
data = DataLoader(train_dataset, batch_size=batch_size)
optimizer = torch.optim.Adam(model.parameters())

#funkcja kosztu
criterion = torch.nn.NLLLoss()

model.train()
step = 0
for x, y in data:
   x = x.to(device)
   y = y.to(device)
   optimizer.zero_grad()
   ypredicted = model(x)
   loss = criterion(torch.log(ypredicted), y)
   if step % 100 == 0:
      print(step, loss)
   step += 1
   loss.backward()
   optimizer.step()

torch.save(model.state_dict(), 'model1.bin')

Evaluation

model = SimpleBigramNeuralLanguageModel(vocab_size, embed_size).to(device)
model.load_state_dict(torch.load('model1.bin'))
model.eval()

execute('./drive/MyDrive/ColabNotebooks/america/dev-0')
execute('./drive/MyDrive/ColabNotebooks/america/test-A')