Rewrite training and inference code to more modern pytorch, add some functionalities and models

app/api/slang.py

@@ -1,24 +1,27 @@
 from flask import jsonify
 from app.api import bp
 
-from app.ml_models.rnn.loaded_rnn_model import return_loaded_model
+from app.ml_models.rnn.loaded_rnn_model import return_loaded_model, load_model
 from app.ml_models.rnn.helpers import random_choice
 from app.ml_models.rnn.generate import generate_word
 
 
 @bp.route("/generate_slang", methods=["GET"])
 def generate_slang():
     """Generate and return a new slang word."""
-    model, ALL_LETTERS = return_loaded_model()
-    N_LETTERS = len(ALL_LETTERS) + 1
 
+    # TODO Should not load model every time a word is queried
+    # I know nothing of flask, can we save the model upon starting the app?
+
+    model, dataset = load_model()
+
+    # TODO: Should check if the word is just memorized
     new_word = generate_word(
         model=model,
-        N_LETTERS=N_LETTERS,
-        ALL_LETTERS=ALL_LETTERS,
-        start_letter=random_choice(ALL_LETTERS),
-        maxn=20,  # TODO: Fix this.
-        temp=0.3,
+        dataset=dataset,
+        start_letter='random',
+        max_len=20,  # TODO: Fix this. (?)
+        temperature=0.3,
     )
 
     # TODO: Return a json containing the word.

app/ml_models/rnn/data_tools.py

@@ -0,0 +1,50 @@
+import torch
+import os
+from torch.utils.data import Dataset, DataLoader
+from collections import Counter
+
+
+class WordLevelDataset(Dataset):
+    def __init__(self,
+                 prefix: str = '../../../data/',
+                 filename_dataset: str = 'straattaal.txt',
+                 filename_vocab: str = 'vocabulary.txt'):
+        filename_dataset = os.path.join(prefix, filename_dataset)
+        filename_vocab = os.path.join(prefix, filename_vocab)
+
+        with open(filename_dataset, 'r', encoding='utf-8') as f:
+            lines = f.read().strip().lower()
+            self.words = [s.strip().replace('\t', '')
+                          for s in lines.split("\n")]
+        with open(filename_vocab, 'r', encoding='utf-8') as f:
+            self.vocabulary = list(f.read())
+        self.vocabulary += ['<BOS>', '<EOS>']
+        self.vocabulary_size = len(self.vocabulary)
+        self.char_to_idx_dict = {ch: i for i, ch in enumerate(self.vocabulary)}
+        self.idx_to_char_dict = {i: ch for i, ch in enumerate(self.vocabulary)}
+
+    def __len__(self):
+        return len(self.words)
+
+    def __getitem__(self, i):
+        s1 = [
+            self.char_to_idx_dict[z]
+            for z in ["<BOS>"] + list(self.words[i])
+        ]
+        s2 = [
+            self.char_to_idx_dict[z]
+            for z in list(self.words[i]) + ["<EOS>"]
+        ]
+        return torch.LongTensor(s1), torch.LongTensor(s2)
+
+    def convert_to_string(self, char_ix):
+        result = "".join(self.idx_to_char_dict[ix] for ix in char_ix)
+        return result
+
+
+if __name__ == "__main__":
+    hi = WordLevelDataset('../../../data/', 'dutch.txt')
+    hi_loader = DataLoader(hi, 1)
+    for z in hi_loader:
+        print(z)
+        break

app/ml_models/rnn/generate.py

@@ -1,54 +1,74 @@
 import torch
 from app.ml_models.rnn.helpers import get_input_tensor
 from app.ml_models.rnn.rnn_model import RNN
+from random import choice as choose
 
 
-def generate_word(
-    model: RNN,
-    N_LETTERS: int,
-    ALL_LETTERS: set,
-    start_letter: str = "a",
-    maxn: int = 20,
-    temp: float = 0,
-) -> str:
-    """Generate a new Slang word.
+def next_char(out, temperature):
+    # Softmax of the last dimension
+    if torch.distributions.Uniform(0, 1).sample() < temperature:
+        probs = torch.softmax((out), -1)
+        #probs = torch.softmax(temperature*(out), -1) # This is good for randomness (temperature < 1)
+        choice = torch.multinomial(probs.squeeze(0), 1)
+    else:
+        choice = torch.argmax(out, dim=2)
+    return choice
+
+
+def generate_word(model, dataset, start_letter=None,  max_len=20, temperature=0.25, device='cpu'):
+    """Generate a new word.
 
     Args:
         model: Pre-trained Recurrent Neural Network model.
-        N_LETTERS: Number of unique letters found in the training corpus.
-        ALL_LETTERS: All letters found in the training data.
+        dataset: WordLevelDataset object 
         start_letter: Letter to start the word with.
-        maxn: Maximum number of letters to be used.
+        max_len: Maximum number of letters to be used.
         temp: Temperature used for sampling.
+        device: torch device string
     """
+    # Evaluation mode
+    model.eval()
+    # no gradient
     with torch.no_grad():
-        input = get_input_tensor(
-            start_letter,
-            N_LETTERS=N_LETTERS,
-            ALL_LETTERS=ALL_LETTERS,
-        )
-        hidden = model.initHidden()
-        output_name = start_letter
-        for i in range(maxn):
-            output, hidden = model(input[0], hidden)
-            if torch.distributions.Uniform(0, 1).sample() < temp:
-                probs = torch.softmax(output, 1)
-                dist = torch.distributions.Categorical(probs)
-                pick = dist.sample()
-            else:
-                topv, topi = output.topk(1)
-                pick = topi[0][0]
-            if pick == N_LETTERS - 1:
-                break
-            else:
-                letter = ALL_LETTERS[pick]
-                output_name += letter
-            input = get_input_tensor(
-                letter,
-                N_LETTERS=N_LETTERS,
-                ALL_LETTERS=ALL_LETTERS)
-
-        return output_name
+        # Hidden stuff initialized to None (pytorch makes this zeros automatically)
+        h = None
+        it = 0
+
+        # Always generate the Beginning of Word token first and feed it to the RNN
+        # TODO: Maybe don't ? It results in a lot of copying behaviour for small datasets
+        # idxs = torch.Tensor([dataset.char_to_idx_dict["<BOS>"]]
+        #                     ).long().unsqueeze(0).to(device)
+        # out, h = model(idxs, h)
+
+        choice = torch.Tensor([-99])
+
+        # Generate a random choice from the vocabulary and put it in the to-be-fed IDXs
+        if start_letter == 'random':
+            letters_idx = torch.Tensor(
+                [dataset.char_to_idx_dict[choose("abcdefghijklmnopqrstuvwxyz")]]
+            ).long().unsqueeze(0).to(device)
+
+        # Generate a random choice from the input
+        elif start_letter is not None:
+            letters_idx = torch.Tensor(
+                [dataset.char_to_idx_dict[choose(start_letter)]]
+            ).long().unsqueeze(0).to(device)
+
+        # Let the RNN decide for this first round.
+        else:
+            choice = next_char(out, temperature)
+            letters_idx = choice.to(device)
+
+        # Check if the token is an EOS token.
+        while choice.item() != dataset.char_to_idx_dict["<EOS>"] and it < max_len:
+            # Pass the latest character to the model, store new hidden stuff.
+            out, h = model(letters_idx[it:], h)
+            choice = next_char(out, temperature)
+            letters_idx = torch.cat((letters_idx, choice), 0)
+            it += 1
+
+        output_string = letters_idx.squeeze(1).tolist()
+    return dataset.convert_to_string(output_string).split('<EOS>')[0]
 
 
 # # TODO: Move this piece of code to generate.py?

app/ml_models/rnn/loaded_rnn_model.py

@@ -1,6 +1,25 @@
 import os
 import torch
-from app.ml_models.rnn.rnn_model import RNN
+from app.ml_models.rnn.rnn_model import RNN, RNNAnna
+from app.ml_models.rnn.data_tools import WordLevelDataset
+
+
+def load_model(model_name: str = '2021_straattaal_epoch100.pt', device: str = 'cpu'):
+    """
+    Args
+        model_name: Filename of the model
+        device: CUDA device name to map to, probably cpu
+    """
+    path = os.path.join(os.path.abspath(os.getcwd()),
+                        "app", "ml_models", "rnn")
+    path = os.path.join(path, model_name)
+    dataset = WordLevelDataset('data/', 'straattaal.txt')
+    # TODO: Fix hardcoded hidden size
+    m = RNNAnna(dataset.vocabulary_size, 128)
+    checkpoint = torch.load(path, map_location=torch.device(device))
+    m.load_state_dict(checkpoint['model_state_dict'])
+    m.eval()
+    return m, dataset
 
 
 def return_loaded_model():

app/ml_models/rnn/rnn_model.py

@@ -24,3 +24,26 @@ def forward(self, input, hidden):
 
     def initHidden(self):
         return torch.zeros(1, self.hidden_size)
+
+
+class RNNAnna(nn.Module):
+    def __init__(
+        self,
+        vocab_size,
+        hidden_size,
+        train_embeddings=False,
+    ):
+        super(RNNAnna, self).__init__()
+        self._embedding = nn.Embedding(vocab_size, vocab_size)
+        self._embedding.weight.data = torch.eye(vocab_size)
+        self._embedding.weight.requires_grad = train_embeddings
+
+        self.lstm = nn.RNN(vocab_size, hidden_size, 1, batch_first=False)
+        self.dropout = nn.Dropout(0.1)
+        self.final = nn.Linear(hidden_size, vocab_size)
+        self.hidden_size = hidden_size
+
+    def forward(self, x, hidden=None):
+        x = self._embedding(x)
+        out, hidden = self.lstm(x, hidden)
+        return self.final(self.dropout(out)), hidden

app/ml_models/rnn/train.py

@@ -0,0 +1,82 @@
+import torch.nn as nn
+import torch
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from app.ml_models.rnn.data_tools import WordLevelDataset
+from app.ml_models.rnn.rnn_model import RNNAnna
+from app.ml_models.rnn.generate import generate_word
+
+
+def train(rnn,
+          dataloader,
+          dataset,
+          learning_rate=0.0005,
+          epochs=500,
+          device='cpu',
+          name='straattaal',
+          save_every=50,
+          print_every=10000):
+
+    # With CrossEntropyLoss we don't need (manual) one-hot
+    criterion = nn.CrossEntropyLoss()
+
+    # Use SGD optimizer so we don't need manual param updates.
+    optimizer = torch.optim.SGD(
+        rnn.parameters(), lr=learning_rate, momentum=0.9)
+    for epoch in range(epochs):
+        total_loss = 0
+        rnn.train()
+        for i, (input_line_tensor, target_line_tensor) in tqdm(enumerate(dataloader), total=len(dataloader)):
+            optimizer.zero_grad()
+            input_line_tensor = input_line_tensor.to(device)
+            target_line_tensor = target_line_tensor.to(device)
+
+            # Run model ye olde way
+            #output, _  = rnn(input_line_tensor)
+            #loss = criterion(output.permute(1, 2, 0), target_line_tensor.permute(1,0))
+
+            # Run model ye new way
+            loss = 0
+            hidden = None
+            for Z in range(input_line_tensor.size(1)):
+                # TODO unsqueeze is necessary for batch size 1
+                # Make this generic for larger batch size (it will also be faster on bigger dataset)
+                output, hidden = rnn(
+                    input_line_tensor[:, Z].unsqueeze(1), hidden)
+                l = criterion(output.permute(
+                    1, 2, 0), target_line_tensor[:, Z].unsqueeze(1).permute(1, 0))
+                loss += l
+
+            loss.backward()
+            optimizer.step()
+            total_loss += loss.item()
+            if (i+1) % print_every == 0:
+                print('Loss', total_loss / i)
+                for _ in range(10):
+                    print('\t', generate_word(
+                        rnn, dataset, start_letter="afhklmnopqrstu", temperature=0.3, device=device))
+                rnn.train()
+
+        # TODO plot loss... maybe.... store it somewhere.... im too lazy
+        if epoch % save_every == 0:
+            print('Loss', total_loss / i)
+            for _ in range(10):
+                print('\t', generate_word(
+                    rnn, dataset, start_letter="abcdefghijklmnoprstuvwz", temperature=0.3, device=device))
+
+            # TODO Save this to some generic spot, not just aat cwd...
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': rnn.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict()
+            }, f"{name}_statedict_{epoch}.pt")
+
+
+if __name__ == "__main__":
+    hi = WordLevelDataset('../../../data/', 'straattaal.txt')
+
+    # Currently only batch size 1 works
+    hi_loader = DataLoader(hi, 1, shuffle=True)
+    rnn = RNNAnna(hi.vocabulary_size, 64, 128)
+    train(rnn, hi_loader, hi)