Upload model.py

model.py

@@ -0,0 +1,175 @@
+from dataclasses import dataclass
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+import math
+
+def precompute(head_dim : int, seq_len : int, device : str, base : float = 10000.0):
+    assert head_dim % 2 == 0
+
+    theta_pair = torch.arange(0,head_dim,2).float()
+    theta = 1.0 / (base ** (theta_pair / head_dim)).to(device)
+
+    m = torch.arange(seq_len,device=device)
+    freqs = torch.outer(m,theta).float()
+    
+    freqs_complex = torch.polar(torch.ones_like(freqs),freqs)
+
+    return freqs_complex
+
+def apply_rope(x : torch.tensor, freqs_complex : torch.tensor, device:str):
+    x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1],-1,2))
+    freqs_complex = freqs_complex.unsqueeze(0).unsqueeze(0)
+
+    x_rotated = x_complex * freqs_complex
+    x_out = torch.view_as_real(x_rotated)
+    x_out = x_out.reshape(*x.shape)
+
+    return x_out.type_as(x)
+
+
+class CasualSelfAttention(nn.Module):
+    def __init__(self,config):
+        super().__init__()
+
+        self.c_attn = nn.Linear(config.n_embd, config.n_embd * 3,bias=False)
+        self.c_proj = nn.Linear(config.n_embd,config.n_embd,bias=False)
+
+        self.n_head = config.n_head
+        self.n_embd = config.n_embd
+
+
+    def forward(self,x,freqs):
+        B,T,C = x.size()
+        qkv = self.c_attn(x)
+        q,k,v = qkv.split(self.n_embd,dim=2)
+
+        q = q.view(B,T, self.n_head, C // self.n_head).transpose(1,2)
+        k = k.view(B,T, self.n_head, C // self.n_head).transpose(1,2)
+        v = v.view(B,T , self.n_head, C // self.n_head).transpose(1,2)
+
+        q = apply_rope(q, freqs, x.device)
+        k = apply_rope(k, freqs, x.device)
+
+        y = F.scaled_dot_product_attention(q,k,v,is_causal=True)
+        y = y.transpose(1,2).contiguous().view(B,T,C)
+
+        y = self.c_proj(y)
+
+
+        return y
+    
+
+
+class FeedForward(nn.Module):
+    def __init__(self,config):
+        super().__init__()
+        hidden_dim = int(2 * config.n_embd / 3)
+        hidden_dim = config.multiple_of * ((hidden_dim + config.multiple_of - 1) // config.multiple_of)
+
+        self.w1 = nn.Linear(config.n_embd,hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim,config.n_embd, bias = False)
+        self.w3  = nn.Linear(config.n_embd, hidden_dim, bias= False)
+
+    def forward(self,x : torch.tensor):
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+    
+
+class Block(nn.Module):
+    def __init__(self,config):
+        super().__init__()
+
+        self.sa = CasualSelfAttention(config)
+        self.mlp = FeedForward(config)
+        self.ln1 = nn.RMSNorm(config.n_embd,eps=1e-05)
+        self.ln2 = nn.RMSNorm(config.n_embd,eps =1e-05)
+
+
+    def forward(self,x,freqs):
+        x = x + self.sa(self.ln1(x),freqs)
+        x = x + self.mlp(self.ln2(x))
+
+        return x
+    
+
+class Model(nn.Module):
+    def __init__(self,config):
+        super().__init__()
+
+        self.register_buffer(
+            'freqs_complex',precompute(
+                config.n_embd // config.n_head,
+                config.block_size,
+                device=config.device
+            )
+        )
+
+        self.config = config
+        self.transformer = nn.ModuleDict(dict(
+            wte = nn.Embedding(config.vocab_size,config.n_embd),
+            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
+            ln_f = nn.RMSNorm(config.n_embd,eps = 1e-05)
+        ))
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size,bias=False)
+
+
+        self.transformer.wte.weight = self.lm_head.weight
+
+    def forward(self,idx,targets=None):
+        B,T = idx.size()
+
+        x = self.transformer.wte(idx)
+
+        for block in self.transformer.h:
+            x = block(x,self.freqs_complex[:T])
+
+        x = self.transformer.ln_f(x)
+        logits = self.lm_head(x)
+
+
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(logits.view(-1,logits.size(-1)),targets.view(-1))
+
+        return logits,loss
+
+
+    def sample_top_p(self,probs,p):
+        probs_sort, probs_idx = torch.sort(probs,dim=-1,descending=True)
+        probs_sum = torch.cumsum(probs_sort,dim=-1)
+        mask = probs_sum - probs_sort > p
+        probs_sort[mask] = 0.0
+        probs_sort.div_(probs_sort.sum(dim=-1,keepdim=True))
+
+        next_token = torch.multinomial(probs_sort,num_samples=1)
+        next_token = torch.gather(probs_idx,-1,next_token)
+
+        return next_token
+
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None,top_p = None):
+
+        for _ in range(max_new_tokens):
+
+            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
+
+            logits, _ = self(idx_cond)
+
+            logits = logits[:, -1, :] / temperature
+
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float('Inf')
+                
+
+            probs = F.softmax(logits, dim=-1)
+
+            
+            if top_p is not None:
+                idx_next = self.sample_top_p(probs,top_p)
+            else:
+                idx_next = torch.multinomial(probs, num_samples=1)
+
+            idx = torch.cat((idx, idx_next), dim=1)
+
+        return idx