Finalized project

.gitignore

@@ -211,7 +211,7 @@ dmypy.json
 
 ### VisualStudioCode ###
 .vscode/*
-!.vscode/settings.json
+.vscode/settings.json
 !.vscode/tasks.json
 !.vscode/launch.json
 !.vscode/extensions.json
@@ -229,8 +229,11 @@ model
 train/runs
 train/output*
 train/models
-
+train/ru_gpt2
 # prep files
 data/prep
 # jokes database
-jokes.db
+!jokes.db
+*.db
+# log files
+*.log

bot/bot.cfg

@@ -3,6 +3,7 @@ token = ...
 ab_test = false
 [model]
 model_paths = model
+rus_model_path = rus_model
 dataset_paths = data/qa_jokes.csv
 max_joke_len = 40
 buffer_size = 16

bot/data.py

@@ -0,0 +1,38 @@
+import os
+import logging
+import pandas as pd
+
+
+class Dataset:
+    """Wrapper for the DataFrame to return values similar 
+    to `AbstractJokeGenerator` output.
+    """
+
+    def __init__(self, dataset_path, promt_token, answer_token):
+        self.promt_token = promt_token
+        self.answer_token = answer_token
+        self.name = os.path.split(dataset_path)[1]
+        self.data = pd.read_csv(dataset_path)
+        self.logger = logging.getLogger("DS: " + self.name)
+
+    def __getitem__(self, idx):
+        question = self.data['Question'].iloc[idx].strip()
+        answer = self.data['Answer'].iloc[idx].strip()
+        text = (self.promt_token + question + '\n'
+                + self.answer_token + ' ' + answer)
+        self.logger.info('Got joke from dataset')
+        return {
+            'text': text,
+            'generated_by': self.name,
+        }
+
+    def __len__(self):
+        return len(self.data)
+
+
+class Joke:
+    """An interface class for a Joke."""
+
+    def __init__(self, text, id):
+        self.id = id
+        self.text = text

bot/inference.py

@@ -1,4 +1,6 @@
 import torch
+from torch import LongTensor
+import logging
 
 from transformers import (
     CTRLLMHeadModel,
@@ -14,48 +16,57 @@
     XLNetLMHeadModel,
     XLNetTokenizer,
 )
+from yt_encoder import YTEncoder
 
 MODEL_CLASSES = {
     "gpt2": (GPT2LMHeadModel, GPT2Tokenizer),
+    "gpt2-yttm": (GPT2LMHeadModel, YTEncoder),
     "ctrl": (CTRLLMHeadModel, CTRLTokenizer),
     "openai-gpt": (OpenAIGPTLMHeadModel, OpenAIGPTTokenizer),
     "xlnet": (XLNetLMHeadModel, XLNetTokenizer),
     "transfo-xl": (TransfoXLLMHeadModel, TransfoXLTokenizer),
     "xlm": (XLMWithLMHeadModel, XLMTokenizer),
 }
 
+# Don't show warnings.
+logging.getLogger("transformers.tokenization_utils").setLevel(logging.ERROR)
+logging.getLogger("transformers.modeling_utils").setLevel(logging.ERROR)
+logging.getLogger("transformers.configuration_utils").setLevel(logging.ERROR)
+
 
 class ModelWrapper:
     def __init__(self, model_path, model_name,
                  device='cpu',
                  model_type='gpt2',
-                 max_length=40,
+                 max_len=40,
                  temperature=0.9,
-                 num_return_sequences=1,
+                 n_return_sequences=1,
                  repetition_penalty=1.0,
                  k=50,
-                 p=0.95,):
-        self.num_return_sequences = num_return_sequences
+                 p=0.95,
+                 **kwargs):
+        self.n_return_sequences = n_return_sequences
         self.repetition_penalty = repetition_penalty
         self.k = k
         self.p = p
         self.temperature = temperature
-        self.device = device = torch.device(device)
-        self.max_length = max_length
+        self.device = torch.device(device)
+        self.max_length = max_len
         model_class, tokenizer_class = MODEL_CLASSES[model_type]
         self.tokenizer = tokenizer_class.from_pretrained(model_path)
         self.model = model_class.from_pretrained(model_path)
         self.name = model_name
-        self.model.to(device)
+        self.logger = logging.getLogger('ML: ' + self.name)
+        self.model.to(self.device)
 
     def __encode(self, text):
-        encoded_prompt = self.tokenizer.encode(
-            text, add_special_tokens=False, return_tensors="pt")
+        encoded_prompt = self.tokenizer.encode(text, add_special_tokens=False)
+        encoded_prompt = LongTensor(encoded_prompt).unsqueeze(0)
         return encoded_prompt.to(self.device)
 
-    def generate(self, beginning, num_return_sequences=None):
-        if num_return_sequences is None:
-            num_return_sequences = self.num_return_sequences
+    def generate(self, beginning, n_return_sequences=None):
+        if n_return_sequences is None:
+            n_return_sequences = self.n_return_sequences
         encoded_prompt = self.__encode(beginning)
         output_sequences = self.model.generate(
             input_ids=encoded_prompt,
@@ -65,7 +76,7 @@ def generate(self, beginning, num_return_sequences=None):
             top_p=self.p,
             repetition_penalty=self.repetition_penalty,
             do_sample=True,
-            num_return_sequences=num_return_sequences,
+            num_return_sequences=n_return_sequences,
         )
 
         # Remove the batch dimension when returning multiple sequences
@@ -86,7 +97,7 @@ def generate(self, beginning, num_return_sequences=None):
     #     dt = datetime.datetime.now() - start
     #     print("\t", dt)
     # Batch processing test
-    res = m.generate("[QUESTION] ", num_return_sequences=4)
+    res = m.generate("[QUESTION] ", n_return_sequences=4)
     for j in res:
         print(j)