| 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 |
