Create model.py

model.py

@@ -0,0 +1,249 @@
+import math
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from utils import DEVICE
+
+class RMSNorm(torch.nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = self._norm(x.float()).type_as(x)
+        return output * self.weight
+
+
+class Attention(nn.Module):
+    """
+    Multi-head Self-Attention with RoPE
+    """
+
+    def __init__(self, num_heads, head_size, num_embed, dropout):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_size = head_size
+
+        self.wq = nn.Linear(num_embed, num_heads * head_size, bias = False)
+        self.wk = nn.Linear(num_embed, num_heads * head_size, bias = False)
+        self.wv = nn.Linear(num_embed, num_heads * head_size, bias = False)
+        self.wo = nn.Linear(num_heads * head_size, num_embed, bias = False)
+
+        inv_freq = 1 / (500000 ** (torch.arange(0, head_size, 2)[: (head_size // 2)].float() / head_size))
+        self.register_buffer('inv_freq', inv_freq)
+
+        self.dropout = nn.Dropout(dropout)
+
+    def reshape_for_broadcast(self, freq_cis, x):
+        ndim = x.ndim
+        shape = [1] * (ndim - 2) + list(freq_cis.shape)
+        return freq_cis.view(*shape)
+
+    def apply_rope(self, x, position, freq):
+        t = torch.arange(position, device=freq.device, dtype=torch.float32)
+        freq = torch.outer(t, freq)
+        freq_cis = torch.polar(torch.ones_like(freq), freq)
+        x_ = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+        freq_cis = self.reshape_for_broadcast(freq_cis, x)
+        x_out = torch.view_as_real(x_ * freq_cis).flatten(3)
+        return x_out.type_as(x)
+
+    def forward(self, x):
+        B, T, C = x.shape
+
+        mask = torch.triu(torch.full((T, T), float("-inf"), device=x.device), diagonal=1)
+
+        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
+
+        xq = xq.view(B, T, self.num_heads, self.head_size)
+        xk = xk.view(B, T, self.num_heads, self.head_size)
+        xv = xv.view(B, T, self.num_heads, self.head_size)
+
+        xq = xq.transpose(1, 2)
+        xk = xk.transpose(1, 2)
+        xv = xv.transpose(1, 2)
+
+        xq = self.apply_rope(xq, T, self.inv_freq)
+        xk = self.apply_rope(xk, T, self.inv_freq)
+
+        attn_weights = torch.matmul(xq, xk.transpose(2, 3)) / math.sqrt(self.head_size)
+        attn_weights += mask
+        attn_weights = F.softmax(attn_weights.float(), dim=-1).type_as(xq)
+        output = torch.matmul(attn_weights, xv)
+        output = output.transpose(1, 2).contiguous().view(B, T, C)
+        return self.dropout(self.wo(output))
+
+
+class MLP(nn.Module):
+    """
+    Implementation of a Multi-Layer Perceptron (MLP) sub-module.
+
+    This module is a simple feed-forward network with two hidden layers
+    used in various Transformer components like the Mixture of Experts layer.
+    """
+
+    def __init__(self, num_embed, dropout):
+        """
+        Constructor for the MLP.
+
+        Args:
+        num_embed (int): The number of embedding dimensions.
+        """
+
+        super().__init__()
+        hidden = int(4 * num_embed * 2 / 3)
+
+        # Define linear layers for the MLP
+        self.w1 = nn.Linear(num_embed, hidden, bias=False)
+        self.w2 = nn.Linear(hidden, num_embed, bias=False)
+
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        """
+        Forward pass of the MLP.
+
+        Args:
+        x (torch.Tensor): Input tensor of shape (batch_size, seq_len, num_embed).
+
+        Returns:
+        torch.Tensor: Output tensor after passing through the MLP (shape: batch_size, seq_len, num_embed).
+        """
+        return self.dropout(self.w2(F.silu(self.w1(x))))
+
+class TransformerBlock(nn.Module):
+    """
+    This calss will group together MultiHead Attention and
+    MLP, so that we can copy it in Transformer
+    """
+
+    def __init__(self, num_heads, head_size, num_embed, dropout):
+        super().__init__()
+
+        self.mha = Attention(
+            num_heads=num_heads,
+            head_size=head_size,
+            num_embed=num_embed,
+            dropout=dropout
+        )
+
+        self.mlp = MLP(num_embed = num_embed, dropout = dropout)
+
+        # add the layer normalization
+        self.norm1 = RMSNorm(num_embed)
+        self.norm2 = RMSNorm(num_embed)
+
+    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).
+        """
+        #print(x.shape)
+        x = x + self.mha(self.norm1(x))
+        x = x + self.mlp(self.norm2(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.model_type = 'Prome'
+        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.head_size = kwargs.get("head_size", 128)
+        self.num_layers = kwargs.get("num_layers", 4)
+        self.dropout = kwargs.get("dropout", 0.2)
+        self.max_seq_len = kwargs.get("max_sqe_len", 1024)
+        # each token reads the logits for the next token from a lookup table
+        self.token_embedding_table = nn.Embedding(self.vocab_size, self.num_embed)
+        # each position from 0 to block_size-1 will get its embedding
+        #self.position_embedding_table = nn.Embedding(self.max_seq_len, self.num_embed)
+
+        self.decoder = nn.Sequential(
+            *[
+                TransformerBlock(
+                    num_heads=self.num_heads,
+                    head_size=self.head_size,
+                    num_embed=self.num_embed,
+                    dropout=self.dropout,
+                )
+                for _ in range(self.num_layers)
+            ]
+        )
+
+        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
+        x = self.token_embedding_table(idx)
+        # (T, C)
+        #posit_emb = self.position_embedding_table(torch.arange(T, device=DEVICE))
+
+        #x = token_emb + posit_emb
+
+        x = self.decoder(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
+            #logits = logits.to(dtype=torch.float32)
+
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
+        else:
+            loss = None
+
+        return logits, loss
+
+    def generate(self, idx: torch.Tensor, max_new_tokens: int, temperature: float = 0.6, top_p: float = 0.9):
+        for _ in range(max_new_tokens):
+            idx_crop = idx[:, -self.max_seq_len:]
+
+            logits, loss = self.forward(idx_crop)
+            logits = logits[:, -1, :]
+
+            if temperature > 0:
+                probs = F.softmax(logits / temperature, dim=-1)
+                idx_next = self.sample_top_p(probs, top_p)
+            else:
+                probs = F.softmax(logits, dim=-1)
+                idx_next = torch.multinomial(probs, num_samples=1)
+            idx = torch.cat((idx, idx_next), dim=1)  # (B, T+1)
+        return idx
+
+    def sample_top_p(self, probs: torch.Tensor, top_p: float) -> torch.Tensor:
+        sorted_probs, sorted_indices = torch.sort(probs, descending=True, dim=-1)
+        cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
+
+        # Create a mask for top-p filtering
+        top_p_mask = cumulative_probs <= top_p
+        top_p_mask[..., 1:] = top_p_mask[..., :-1].clone()
+        top_p_mask[..., 0] = 1
+
+        filtered_probs = sorted_probs * top_p_mask
+        filtered_probs /= filtered_probs.sum(dim=-1, keepdim=True)  # Normalize filtered probabilities
+
+        next_token = torch.multinomial(filtered_probs, num_samples=1)
+        return torch.gather(sorted_indices, -1, next_token)