s430705

run2.py

@@ -0,0 +1,187 @@
+import csv
+import itertools
+from os.path import exists
+
+import pandas as pd
+import regex as re
+import torch
+from torch import nn
+from torch.utils.data import DataLoader
+from torchtext.vocab import build_vocab_from_iterator
+
+IN_INPUT_PATH = "train/in.tsv.xz"
+IN_OUTPUT_PATH = "train/expected.tsv"
+VOCAB_SIZE = 30000
+EMBED_SIZE = 150
+BATCH_SIZE = 8000
+DEV_PATH = "dev-0/"
+TEST_PATH = "test-A/"
+DEVICE = "cpu"
+
+
+def clean(text):
+    text = str(text).lower().replace("-\\n", "").replace("\\n", " ")
+    return re.sub(r"\p{P}", "", text)
+
+
+def get_words_from_line(line, specials=True):
+    line = line.rstrip()
+    if specials:
+        yield "<s>"
+    for m in re.finditer(r"[\p{L}0-9\*]+|\p{P}+", line):
+        yield m.group(0).lower()
+    if specials:
+        yield "</s>"
+
+
+def get_word_lines_from_data(d):
+    for line in d:
+        yield get_words_from_line(line)
+
+
+def look_ahead_iterator(gen):
+    w1 = None
+    for item in gen:
+        if w1 is not None:
+            yield (w1, item)
+        w1 = item
+
+
+class SimpleBigramNeuralLanguageModel(torch.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)
+
+
+class Bigrams(torch.utils.data.IterableDataset):
+    def __init__(self, data, vocabulary_size):
+        self.vocab = build_vocab_from_iterator(
+            get_word_lines_from_data(data),
+            max_tokens=vocabulary_size,
+            specials=["<unk>"],
+        )
+        self.vocab.set_default_index(self.vocab["<unk>"])
+        self.vocabulary_size = vocabulary_size
+        self.data = data
+
+    def __iter__(self):
+        return look_ahead_iterator(
+            (
+                self.vocab[t]
+                for t in itertools.chain.from_iterable(
+                    get_word_lines_from_data(self.data)
+                )
+            )
+        )
+
+
+def get_dataset():
+    X_train = pd.read_csv(
+        IN_INPUT_PATH,
+        sep="\t",
+        header=None,
+        quoting=csv.QUOTE_NONE,
+        nrows=200000,
+        on_bad_lines="skip",
+        encoding="UTF-8",
+    )
+    Y_train = pd.read_csv(
+        IN_OUTPUT_PATH,
+        sep="\t",
+        header=None,
+        quoting=csv.QUOTE_NONE,
+        nrows=200000,
+        on_bad_lines="skip",
+        encoding="UTF-8",
+    )
+
+    X_train = X_train[[6, 7]]
+    X_train = pd.concat([X_train, Y_train], axis=1)
+    X_train = X_train[6] + X_train[0] + X_train[7]
+    X_train = X_train.apply(clean)
+    return Bigrams(X_train, VOCAB_SIZE)
+
+
+dataset = get_dataset()
+
+
+def get_model():
+
+    model = SimpleBigramNeuralLanguageModel(VOCAB_SIZE, EMBED_SIZE).to(DEVICE)
+
+    if not exists("nn_model.bin"):
+        data = DataLoader(dataset, batch_size=BATCH_SIZE)
+        optimizer = torch.optim.Adam(model.parameters())
+        criterion = torch.nn.NLLLoss()
+
+        model.train()
+        step = 0
+        for i in range(2):
+            for x, y in data:
+                x = x.to(DEVICE)
+                y = y.to(DEVICE)
+                optimizer.zero_grad()
+                y_predicted = model(x)
+                loss = criterion(torch.log(y_predicted), y)
+                if step % 100 == 0:
+                    print(step, loss)
+                step += 1
+                loss.backward()
+                optimizer.step()
+
+        torch.save(model.state_dict(), "nn_model.bin")
+    else:
+        model.load_state_dict(torch.load("nn_model.bin", map_location=torch.device('cpu')))
+    return model
+
+
+vocab = dataset.vocab
+model = get_model()
+
+
+def predict(ws):
+    ixs = torch.tensor(vocab.forward(ws)).to(DEVICE)
+    out = model(ixs)
+    top = torch.topk(out[0], 8)
+    top_indices = top.indices.tolist()
+    top_probs = top.values.tolist()
+    top_words = vocab.lookup_tokens(top_indices)
+    pred_str = ""
+    for word, prob in list(zip(top_words, top_probs)):
+        pred_str += f"{word}:{prob} "
+    return pred_str
+
+
+def predict_input(file):
+    X_test = pd.read_csv(
+        f"{file}/in.tsv.xz",
+        sep="\t",
+        header=None,
+        quoting=csv.QUOTE_NONE,
+        on_bad_lines="skip",
+        encoding="UTF-8",
+    )[6]
+    X_test = X_test.apply(clean)
+    with open(f"{file}/out.tsv", "w+", encoding="UTF-8") as f:
+        for row in X_test:
+            before = None
+            for before in get_words_from_line(clean(str(row)), False):
+                pass
+            before = [before]
+            if len(before) < 1:
+                pred_str = "a:0.2 the:0.2 to:0.2 of:0.1 and:0.1 of:0.1 :0.1"
+            else:
+                pred_str = predict(before)
+            pred_str = pred_str.strip()
+            f.write(pred_str + "\n")
+
+
+predict_input(DEV_PATH)
+predict_input(TEST_PATH)