Bigram

run.py

@@ -9,15 +9,18 @@ import itertools
 import pandas as pd
 from torch.utils.data import DataLoader
 import csv
+import os
 
 def data_preprocessing(text):
     return re.sub(r'\p{P}', '', text.lower().replace('-\\n', '').replace('\\n', ' ').replace("'ll", " will").replace("-", "").replace("'ve", " have").replace("'s", " is"))
 
-def get_words_from_line(line):
+def get_words_from_line(line, s = True):
   line = line.rstrip()
+  if s:
     yield '<s>'
   for m in re.finditer(r'[\p{L}0-9\*]+|\p{P}+', line):
      yield m.group(0).lower()
+  if s:
     yield '</s>'
 
 
@@ -89,6 +92,8 @@ bigram_data = Bigrams(data, vocab_size)
 
 device = 'cpu'
 model = SimpleBigramNeuralLanguageModel(vocab_size, embed_size).to(device)
+
+if(not os.path.exists('model1.bin')):
     data = DataLoader(bigram_data, batch_size=5000)
     optimizer = torch.optim.Adam(model.parameters())
     criterion = torch.nn.NLLLoss()
@@ -108,6 +113,8 @@ for x, y in data:
         optimizer.step()
 
     torch.save(model.state_dict(), 'model1.bin')
+else:
+    model.load_state_dict(torch.load('model1.bin')) 
 
 vocab = bigram_data.vocab
 prediction = 'the:0.03 be:0.03 to:0.03 of:0.025 and:0.025 a:0.025 in:0.020 that:0.020 have:0.015 I:0.010 it:0.010 for:0.010 not:0.010 on:0.010 with:0.010 he:0.010 as:0.010 you:0.010 do:0.010 at:0.010 :0.77'
@@ -131,7 +138,6 @@ def predict(f):
 
     with open(f'{f}/out.tsv', "w+", encoding="UTF-8") as f:
         for row in x:
-            result = {}
             before = None
             for before in get_words_from_line(data_preprocessing(str(row)), False):
                 pass
@@ -144,5 +150,5 @@ def predict(f):
             pred_str = pred_str.strip()
             f.write(pred_str + "\n")
 
-prediction("dev-0/")
-prediction("test-A/")
+predict("dev-0/")
+predict("test-A/")