Delete model.py

model.py

@@ -1,359 +0,0 @@
-import torch
-import torch.nn as nn
-from torch.nn import functional as F
-from utils import DEVICE
-
-class PromeLayerNorm(nn.Module):
-    def __init__(self, epsilon=1e-5):
-        super().__init__()
-        self.epsilon = epsilon
-
-    def forward(self, x):
-        g = torch.nn.Parameter(torch.ones(x.shape[-1])).to(x.device)
-        b = torch.nn.Parameter(torch.zeros(x.shape[-1])).to(x.device)
-
-        u = x.mean(-1, keepdim=True)
-        s = (x - u).pow(2).mean(-1, keepdim=True)
-        x = (x - u) * torch.rsqrt(s + self.epsilon)
-        x = x * g + b
-
-        return x
-
-class PromeStand(nn.Module):
-    def __init__(self, epsilon=1e-5):
-        super().__init__()
-        self.epsilon = epsilon
-
-    def forward(self, x):
-        """
-        x: Input tensor
-        """
-        mean = x.mean() + self.epsilon
-        std = x.std() + self.epsilon
-        x = x - mean
-        x = x / std
-        return x
-
-class PromeEmbedding(nn.Module):
-    """
-    This class implements a Prome embedding layer.
-
-    Args:
-        vocab_size (int): The size of the vocabulary.
-        embedding_dim (int): The dimension of the embedding.
-        padding_idx (int, optional): The padding index. If this is not None, then the padding index will be masked out when calculating the embedding.
-
-    Returns:
-        torch.Tensor: A tensor of shape (batch_size, sequence_length, embedding_dim).
-    """
-    def __init__(self, vocab_size, embedding_dim, padding_idx = None):
-        super().__init__()
-        self.embedding_dim = embedding_dim
-        self.weight = torch.nn.Parameter(torch.randn(vocab_size, embedding_dim))
-        self.padding_idx = padding_idx
-        self.context_matrix = torch.nn.Parameter(torch.randn(vocab_size, embedding_dim))
-
-    def forward(self, input_ids):
-        """
-        Calculates the embedding for the given input IDs.
-
-        Args:
-            input_ids (torch.Tensor): A tensor of shape (batch_size, sequence_length).
-
-        Returns:
-            torch.Tensor: A tensor of shape (batch_size, sequence_length, embedding_dim).
-        """
-        input_ids = input_ids.long()
-        if self.padding_idx is not None:
-            input_ids = input_ids.masked_fill(input_ids == self.padding_idx, 0)
-
-        # get symbol vector
-        embeddings = self.weight[input_ids]
-
-        # Dynamically update context vector based on input embeddings
-        context_vectors = self.context_matrix[input_ids]
-
-        # Modify embeddings using context vector
-        output = embeddings + context_vectors
-
-        return output
-
-class AttentionHead(nn.Module):
-    """
-    One head of the self-attention layer
-    """
-
-    def __init__(self, head_size, num_embed, block_size, dropout):
-        super().__init__()
-        self.key = nn.Linear(num_embed, head_size, bias=False)
-        self.query = nn.Linear(num_embed, head_size, bias=False)
-        self.value = nn.Linear(num_embed, head_size, bias=False)
-        # tril is a lower triangular matrix. it is not a parameter
-        # of the model, so we assign it to the module using register_buffer
-        self.register_buffer("tril", torch.tril(torch.ones(block_size, block_size)))
-
-        # layer norm
-        self.norm = PromeStand()
-
-        # Dropout
-        self.dropout = nn.Dropout(dropout)
-
-
-    def forward(self, x):
-        B, T, C = x.shape
-        key = self.key(x)
-        query = self.query(x)
-        # compute attention scores
-        # (B, T, C) @ (B, C, T) -> (B, T, T)
-        wei = (query @ key.transpose(-2, -1)) * C ** -0.5
-        # Tril matrix (lower triagular matrix) is used to mask 
-        # future positions (setting them to -inf) so that the
-        # decoder "learns" to predict next words
-        wei = wei.masked_fill(self.tril[:T, :T] == 0, -float("inf"))  # (B,T,T)
-        wei = F.silu(F.softmax(wei, dim=-1))
-        # scale
-        # multiplicative attention
-        score = -1 / (C ** -0.5)
-        wei.mul_(score)
-        # weighted aggregation of the values
-        value = self.value(x)
-        out = wei @ value  # (B,T,T) @ (B,T,C) ---> (B,T,C)
-
-        return out
-    
-class MultiHeadAttention(nn.Module):
-    """
-    Multiple Heads of self-attention in parallel
-    """
-
-    def __init__(self, num_heads, head_size, num_embed, block_size, dropout):
-        super().__init__()
-        self.heads = nn.ModuleList(
-            [
-                AttentionHead(
-                    head_size=head_size,
-                    num_embed=num_embed,
-                    block_size=block_size,
-                    dropout=dropout
-                )
-                for _ in range(num_heads)
-            ]
-        )
-        self.proj = nn.Linear(num_embed, num_embed)
-        self.dropout = nn.Dropout(dropout)
-        self.norm = PromeStand()
-
-    def forward(self, x):
-        # output of the self-attention
-        out = torch.concat([h(x) for h in self.heads], dim=-1)
-        # standartization
-        out = self.norm(out + x)
-        # apply the linear projection layer
-        out = self.dropout(self.proj(out))
-
-        return out
-
-
-class MLP(nn.Module):
-    def __init__(self, num_embed, hidden_dim, dropout=0.1):
-        super().__init__()
-        self.dropout = nn.Dropout(dropout)
-        self.fc1 = nn.Linear(num_embed, hidden_dim)
-        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
-        self.fc3 = nn.Linear(hidden_dim, num_embed)
-
-    def forward(self, x):
-        x = self.fc1(x)
-        x = F.silu(x)
-        x = self.fc2(x)
-        x = self.dropout(x)
-        x = F.silu(x)
-        x = self.fc3(x)
-        return x
-
-
-class TransformerBlock(nn.Module):
-    """
-    This calss will group together MultiHead Attention and
-    FeedForward NN, so that we can copy it in Transformer
-    """
-
-    def __init__(self, num_heads, block_size, num_embed, hidden_dim, dropout):
-        super().__init__()
-        head_size = num_embed // num_heads
-        self.mha = MultiHeadAttention(
-            num_heads=num_heads,
-            head_size=head_size,
-            num_embed=num_embed,
-            block_size=block_size,
-            dropout=dropout
-        )
-        self.mlp = MLP(num_embed=num_embed, hidden_dim = hidden_dim, dropout=dropout)
-        # add the layer normalization
-        self.ln = PromeStand(num_embed)
-
-        self.dropout = nn.Dropout(dropout)
-
-    def forward(self, x):
-        """
-        Decodes the input sequence.
-
-        Args:
-            x (torch.Tensor): A tensor of shape (batch_size, sequence_length, embedding_dim).
-            memory (torch.Tensor): A tensor of shape (batch_size, memory_length, embedding_dim).
-
-        Returns:
-            torch.Tensor: A tensor of shape (batch_size, sequence_length, embedding_dim).
-        """
-        y = x
-        
-        x = self.ln(x)
-        x = self.mha(x)
-        x = self.dropout(x)
-        x += y
-        y = x
-        x = self.ln(x)
-        x = self.mlp(x)
-        x = self.mha(x)
-        x += y
-        x = self.dropout(x)
-
-        return x
-    
-
-class TransformerDecoder(nn.Module):
-    """
-    This class implements a Transformer decoder.
-
-    Args:
-        num_heads (int): The number of attention heads.
-        block_size (int): The size of the input sequence.
-        num_embed (int): The dimension of the embedding.
-        num_layers (int): The number of decoder blocks.
-        dropout (float): The dropout rate.
-
-    Returns:
-        torch.Tensor: A tensor of shape (batch_size, sequence_length, embedding_dim).
-    """
-    def __init__(self, num_heads, block_size, num_embed, hidden_dim, num_layers, dropout):
-        super().__init__()
-
-        # Create the embedding layer.
-        self.pemb = PromeEmbedding(block_size, num_embed)
-
-        # Create a sequential block of Transformer blocks.
-        self.blocks = nn.Sequential(
-            *[
-                TransformerBlock(
-                    num_heads=num_heads,
-                    block_size=block_size,
-                    num_embed=num_embed,
-                    hidden_dim = hidden_dim,
-                    dropout=dropout
-                )
-                for _ in range(num_layers)
-            ]
-        )
-
-        # Create a softmax layer.
-        self.softmax = nn.Softmax(dim=-1)
-
-    def forward(self, x):
-        """
-        Decodes the input sequence.
-
-        Args:
-            x (torch.Tensor): A tensor of shape (batch_size, sequence_length).
-
-        Returns:
-            torch.Tensor: A tensor of shape (batch_size, sequence_length, embedding_dim).
-        """
-
-        # Add positional encodings to the input sequence.
-        x = x + self.pemb(torch.arange(x.size(1)))
-
-        x = self.blocks(x)
-
-        # Apply a softmax layer to the output of the last Transformer block.
-        x = self.softmax(x)
-
-        return x
-    
-class Transformer(nn.Module):
-    def __init__(self, **kwargs):
-        super().__init__()
-        # a simple lookup table that stores embeddings of a fixed dictionary and size
-        # each token directly reads off the logits for the next token from a lookup table
-        # see more: https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html
-        self.vocab_size = kwargs.get("vocab_size", 100)
-        self.num_embed = kwargs.get("num_embed", 32)
-        self.block_size = kwargs.get("block_size", 8)
-        self.num_heads = kwargs.get("num_heads", 4)
-        self.num_layers = kwargs.get("num_layers", 4)
-        self.hidden_dim = kwargs.get("hidden_dim", 768)
-        self.dropout = kwargs.get("dropout", 0.2)
-        # each token reads the logits for the next token from a lookup table
-        self.token_embedding_table = PromeEmbedding(self.vocab_size, self.num_embed)
-        # each position from 0 to block_size-1 will get its embedding
-        self.position_embedding_table = PromeEmbedding(self.block_size, self.num_embed)
-
-        self.decoder = TransformerDecoder(self.num_heads, self.block_size, self.num_embed, self.hidden_dim, self.num_layers, self.dropout)
-
-        # we add the layer norm before the Linear layer
-        self.dropout = nn.Dropout(self.dropout)
-        self.ln_f = PromeLayerNorm(self.num_embed)
-        self.lm_head = nn.Linear(self.num_embed, self.vocab_size)
-
-    def forward(self, idx, targets=None):
-        B, T = idx.shape
-        # idx and targets are (B,T) tensor of integers
-        # the token_emb is (B, T, C), C = NUM_EMBED
-        token_emb = self.token_embedding_table(idx)
-        # (T, C)
-        posit_emb = self.position_embedding_table(torch.arange(T, device=DEVICE))
-
-        x = token_emb + posit_emb
-
-        # apply dropout
-        x = self.dropout(x)
-
-        # apply one head of self-attention
-        x = self.decoder(x)
-
-        # apply normalization
-        x = self.ln_f(x)
-
-        # (B, T, vocab_size)
-        logits = self.lm_head(x)
-
-        # Compute the loss
-        if targets != None:
-            # cross_entropy accepts inputs in a (batch_size, num_classes)
-            # so we need to reformat our logits dimensions to
-            # (batch_size * time, dim_vocabulary), time = block_size
-            B, T, C = logits.shape
-            logits = torch.reshape(logits, (B * T, C))
-            targets = torch.reshape(targets, (B * T, ))
-            loss = F.cross_entropy(logits, targets)
-        else:
-            loss = None
-
-        return logits, loss
-
-    def generate(self, idx: torch.Tensor, max_new_tokens: int, block_size: int):
-        # idx is (B, T) array of indices in the current context
-        for _ in range(max_new_tokens):
-            # crop the context too the  last block_size tokens
-            # because tokens don't communicate between blocks
-            idx_crop = idx[:, -block_size:]
-            # get the predictions
-            logits, loss = self.forward(idx_crop)
-            # focus only on the last time step
-            logits = logits[:, -1, :]  # becomes (B, C)
-            # apply softmax to get probabilities
-            probs = F.softmax(logits, dim=-1)  # (B, C)
-            # sample from the distribution with probabilities probs
-            idx_next = torch.multinomial(probs, num_samples=1)  # (B, 1)
-            # append sampled index to the running sequence
-            idx = torch.cat((idx, idx_next), dim=1)  # (B, T+1)
-        return idx