test

zad9.py

@@ -0,0 +1,193 @@
+import itertools
+import lzma
+
+import numpy as np
+import torch
+from torch import nn
+from torch.utils.data import IterableDataset, DataLoader
+from torchtext.vocab import build_vocab_from_iterator
+
+
+def clean_line(line):
+    # Preprocessing
+    separated = line.split('\t')
+    prefix = separated[6].replace(r'\n', ' ')
+    suffix = separated[7].replace(r'\n', ' ')
+    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 n_look_ahead_iterator(n, gen):
+    prevs = [None for _ in range(n)]
+    for item in gen:
+        if prevs[-1] is not None:
+            ngram = prevs[::-1]
+            ngram.append(item)
+            yield np.asarray(ngram)
+        prevs.insert(0, item)
+        prevs = prevs[:n - 1]
+
+
+class Ngrams(IterableDataset):
+    def __init__(self, text_file: str, context_size: int, vocabulary_size: int):
+        self.vocab = build_vocab_from_iterator(
+            get_word_lines_from_file(text_file),
+            max_tokens=vocabulary_size,
+            specials=['<unk>']
+        )
+        self.vocab.set_default_index(self.vocab['<unk>'])
+        self.vocabulary_size = vocabulary_size
+        self.text_file = text_file
+        self.ngram_size = context_size + 1
+
+    def __iter__(self):
+        return n_look_ahead_iterator(
+            self.ngram_size,
+            (self.vocab[t] for t in itertools.chain.from_iterable(get_word_lines_from_file(self.text_file)))
+        )
+
+
+class NgramWithBagLM(nn.Module):
+    def __init__(self, smaller_context_size, context_size, embedding_size, vocabulary_size, hidden_size):
+        super().__init__()
+        self.smaller_context_size = smaller_context_size
+        self.context_size = context_size
+        self.embedding_size = embedding_size
+        self.embedding = nn.Embedding(vocabulary_size, embedding_size)
+        self.bag_embedding = nn.Embedding(vocabulary_size, embedding_size)
+        self.lin1 = nn.Linear((smaller_context_size + 1) * embedding_size, hidden_size)
+        self.rel = nn.ReLU()
+        self.lin2 = nn.Linear(hidden_size, vocabulary_size)
+        self.sm = nn.Softmax(dim=1)
+
+    def forward(self, words):
+        smaller_context_embed = [
+            self.embedding(words[:, i]) for i in range(self.context_size - self.smaller_context_size, self.context_size)
+        ]
+        smaller_context_embed = torch.cat(smaller_context_embed, dim=-1)
+        bag_embed = [
+            self.bag_embedding(words[:, i]) for i in range(self.context_size - self.smaller_context_size)
+        ]
+        bag_embed = torch.mean(torch.stack(bag_embed), dim=0)
+        x = torch.cat((bag_embed, smaller_context_embed), dim=-1)
+
+        x = self.lin1(x)
+        x = self.rel(x)
+        x = self.lin2(x)
+        return self.sm(x)
+
+
+def train_model():
+    model = NgramWithBagLM(
+        smaller_context,
+        context_size,
+        embed_size,
+        vocab_size,
+        hidden_size
+    ).to(device)
+
+    data = DataLoader(train_dataset, batch_size=batch_size)
+    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
+    criterion = torch.nn.NLLLoss()
+
+    model.train()
+
+    step = 0
+    for batch in data:
+        print(batch.shape)
+        # x = batch[:, :context_size]
+        # y = batch[:, context_size]
+        x = batch
+        y = batch[:, left_ctx:left_ctx + 1]
+        print(x.shape)
+        print(y.shape)
+        x = x.to(device)
+        y = y.type(torch.LongTensor)
+        y = y.to(device)
+        optimizer.zero_grad()
+        ypredicted = model(x)
+        loss = criterion(torch.log(ypredicted), y)
+        if torch.isnan(loss):
+            raise Exception("loss is nan")
+        if step % 1000 == 0:
+            print(step, loss)
+        step += 1
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 5)
+        optimizer.step()
+
+    # cond?
+    torch.save(model.state_dict(), path_to_model)
+
+
+def prediction(model, words: list, top=500) -> str:
+    words_tensor = [train_dataset.vocab.forward([word]) for word in words]
+    ixs = torch.tensor(words_tensor).view(-1).to(device)
+    out = model(ixs.view(1, -1))
+    top = torch.topk(out[0], top)
+    top_indices = top.indices.tolist()
+    top_probs = top.values.tolist()
+    top_words = train_dataset.vocab.lookup_tokens(top_indices)
+    zipped = list(zip(top_words, top_probs))
+    for index, element in enumerate(zipped):
+        unk = None
+        if '<unk>' in element:
+            unk = zipped.pop(index)
+            zipped.append(('', unk[1]))
+            break
+    if unk is None:
+        zipped[-1] = ('', zipped[-1][1])
+    return ' '.join([f'{x[0]}:{x[1]}' for x in zipped])
+
+
+device = 'cuda'
+vocab_size = 250
+#context_size = 40
+left_ctx = 20
+right_ctx = 20
+#smaller_context = 8
+smaller_left_ctx = 5
+smaller_right_ctx = 3
+embed_size = 20
+hidden_size = 10
+batch_size = 4000
+lr = 0.0001
+path_to_train = 'train/in.tsv.xz'
+path_to_model = 'model3.bin'
+
+train_dataset = Ngrams(path_to_train, left_ctx + right_ctx, vocab_size)
+
+
+train_model()
+model = NgramWithBagLM(
+        smaller_context,
+        context_size,
+        embed_size,
+        vocab_size,
+        hidden_size
+    ).to(device)
+
+model.load_state_dict(torch.load(path_to_model))
+model.eval()
+
+folder_name = 'dev-0'
+top = 500
+print(f'Creating outputs in {folder_name}')
+with lzma.open(f'{folder_name}/in.tsv.xz', mode='rt', encoding='utf-8') as fid:
+    with open(f'{folder_name}/out-top={top}.tsv', 'w', encoding='utf-8', newline='\n') as f:
+        for line in fid:
+            separated = line.split('\t')
+            prefix = separated[6].replace(r'\n', ' ').split()[-context_size:]
+            output_line = prediction(model, prefix, top)
+            f.write(output_line + '\n')