model.py

import torch

class BigramLanguageModel():
    """
    Class for a bigramlangauemodel implemented from scratch and follow the statistics of the text it been fitted to. 
    """
    def __init__(self, vocab_size, tok_sep=""):
        self.bigramtable = torch.zeros((vocab_size, vocab_size), dtype=torch.long)
        self.vocab_size = vocab_size
        self.token_sep = tok_sep

    def fit(self, text_encoded, alpha=1):
        """
        Function for calculating the bigramtable based on the given encoded text, performing Maximum loglikelhodd + laplace smothing 
        """
        for a, b in zip(text_encoded, text_encoded[1:]):
            self.bigramtable[a, b] += 1

        # Convert to probabilities with laplace smoothing, then row-normalize
        self.bigramtable = (self.bigramtable + alpha).float()
        self.bigramtable /= self.bigramtable.sum(1, keepdim=True)  # each row sums to 1 → conditional distributions

    def save(self, file_path:str):
        """
        Function for saving the bigramtable
        """
        state = {
            "version": 1,
            "vocab_size": self.vocab_size,
            "bigramtable": self.bigramtable,  
            }
        torch.save(state, f"{file_path}.pt")

    def load(self, file_path):
        """
        Function for loading in a bigram model
        """
        state = torch.load(file_path)
        self.bigramtable = state["bigramtable"]
        self.vocab_size = state["vocab_size"]

    def sample(self, bos_id:int, eos_id:int, n=200, temperature=1.0):
        """
        Function to sample from the bigramtable and quit either at 200 characters or when a EOS token appears. 
        """
        if  not isinstance(eos_id, int) or not isinstance(bos_id, int):
             raise Exception(f"The EOS and BOS need to be integers. BOS:{type(bos_id)} EOS:{type(eos_id)}")
        
        x = bos_id
        
        output = [] # don't holds the BOS
        for _ in range(n):
            probs = self.bigramtable[x,:]  # (V,)

            # Temperature scaling to make the gaps bigger/smaller in the distribustion
            if temperature != 1.0:
                logits = torch.log(probs + 1e-12) / temperature
                probs = torch.softmax(logits, dim=0)

            # Smaple a idx in the distribution to generate next token
            x = torch.multinomial(probs, 1).item()
            # Stop if EOS is reached
            if x == eos_id:
                print("Found an EOS.")
                break

            output.append(x)

        return output