model.py

from transformer import *
from training import *


class Model:
    """
    Keeps track of the Adventure state.

    Attributes:
        _vocab: A dataset containing all the vocabulary the network was
            trained on.
        _ai: An instance of a transformer neural network used for text
            generation.
        _text: A list of the last 510 text tokens from the game.
    """

    def __init__(self, model_file, data_file):
        """
        Instanitates game vairables and language model.

        Args:
            model_path: String containing path to file containing pretrained
                model parameters.
        """
        self._vocab = AdventureDataset(f"data/{data_file}").vocab
        self._text = []
        self._ai = Transformer.make_model(len(self._vocab),
                                          len(self._vocab))
        self._ai.load_state_dict(torch.load(f"./model/{model_file}"))

    def generate_text(self, prompt):
        """
        Use the neural network to generate text. Saves the prompt and output.

        Args:
            prompt: String which contains prompt for model to expand apon.
        Returns:
            A string of text generated by the neural network.
        """
        tokenized_prompt = self._vocab.tokenize(prompt)
        self._text += tokenized_prompt
        if len(self._text) > 510:
            self._text = self._text[:len(self._text) - 510]
        output = self._ai.generate_text(self._text, self._vocab)
        tokenized_output = self._vocab.tokenize(output)
        self._text += tokenized_output
        return output


if __name__ == "__main__":
    # Example of the ai's capability
    model = Model("italian_numbers", "test.csv")
    text_in_1 = model._vocab.tokenize("one")
    text_out_1 = model._ai.generate_text(text_in_1, model._vocab)
    text_in_2 = model._vocab.tokenize("two")
    text_out_2 = model._ai.generate_text(text_in_2, model._vocab)
    text_in_3 = model._vocab.tokenize("three")
    text_out_3 = model._ai.generate_text(text_in_3, model._vocab)
    text_in_4 = model._vocab.tokenize("four")
    text_out_4 = model._ai.generate_text(text_in_4, model._vocab)
    print(f"Input: one      Output: {text_out_1}")
    print(f"Input: two      Output: {text_out_2}")
    print(f"Input: three    Output: {text_out_3}")
    print(f"Input: four     Output: {text_out_4}")