import torch from torch.utils.data import Dataset, DataLoader import torch.nn as nn import torch.nn.functional as F class MultiHeadAttention(nn.Module): def __init__(self,d_in,d_out,context_length,dropout,qkv_bias,n_heads): super().__init__() self.n_heads = n_heads self.head_dim = d_out // n_heads self.d_out = d_out self.W_key = nn.Linear(d_in,d_out,bias=qkv_bias) self.W_query = nn.Linear(d_in,d_out,bias=qkv_bias) self.W_value = nn.Linear(d_in,d_out,bias=qkv_bias) self.dropout = nn.Dropout(dropout) self.proj = nn.Linear(d_out,d_out) self.register_buffer( 'mask', torch.triu(torch.ones(context_length, context_length), diagonal=1) ) def forward(self,x): b,n_tokens,d_out = x.shape keys = self.W_key(x).view(b,n_tokens,self.n_heads,self.head_dim) queries = self.W_query(x).view(b,n_tokens,self.n_heads,self.head_dim) values = self.W_value(x).view(b,n_tokens,self.n_heads,self.head_dim) keys = keys.transpose(1,2) queries = queries.transpose(1,2) values = values.transpose(1,2) cntx_vec = F.scaled_dot_product_attention( queries, keys, values, attn_mask=None, dropout_p=self.dropout.p if self.training else 0.0, is_causal=True ) cntx_vec = cntx_vec.transpose(1,2) cntx_vec = cntx_vec.contiguous().view(b,n_tokens,self.d_out) return self.proj(cntx_vec) class NormLayer(nn.Module): def __init__(self,emb_dim): super().__init__() self.eps = 1e-5 self.scale = nn.Parameter(torch.ones(emb_dim)) self.shift = nn.Parameter(torch.zeros(emb_dim)) def forward(self,x): mean = x.mean(dim=-1,keepdim=True) var = x.var(dim=-1,keepdim=True,unbiased=False) return self.scale * ((x-mean)/torch.sqrt(var+self.eps)) + self.shift class GELU(nn.Module): def __init__(self): super().__init__() def forward(self, x): return 0.5 * x * (1 + torch.tanh( torch.sqrt(torch.tensor(2.0 / torch.pi)) * (x + 0.044715 * torch.pow(x, 3)) )) class FeedForward(nn.Module): def __init__(self, cfg): super().__init__() self.layers = nn.Sequential( nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]), GELU(), nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]), ) def forward(self, x): return self.layers(x) class TransformerBlock(nn.Module): def __init__(self,cfg): super().__init__() self.attn = MultiHeadAttention(d_in=cfg["emb_dim"],d_out=cfg["emb_dim"],context_length=cfg["context_length"],dropout=cfg["drop_rate"],qkv_bias=cfg["qkv_bias"],n_heads=cfg["n_heads"]) self.ff = FeedForward(cfg) self.norm1 = NormLayer(cfg["emb_dim"]) self.norm2 = NormLayer(cfg["emb_dim"]) self.drop_shortcut = nn.Dropout(cfg["drop_rate"]) def forward(self,x): shortcut = x x = self.norm1(x) x = self.attn(x) x = self.drop_shortcut(x) x = x + shortcut shortcut = x x = self.norm2(x) x = self.ff(x) x = self.drop_shortcut(x) x = x + shortcut return x vocab_size=50257 class GPTModel(nn.Module): def __init__(self,cfg): super().__init__() self.tok_emb = nn.Embedding(vocab_size,cfg["emb_dim"]) self.pos_emb = nn.Embedding(cfg["context_length"],cfg["emb_dim"]) self.drop_emb = nn.Dropout(cfg["drop_rate"]) self.tranf_blocks = nn.Sequential(*[TransformerBlock(cfg) for _ in range(cfg["n_layers"])]) self.out_head = nn.Linear(cfg["emb_dim"],vocab_size) self.final_norm = NormLayer(cfg["emb_dim"]) def forward(self,x): b,n_inp = x.shape tok_emb = self.tok_emb(x) pos_emb = self.pos_emb(torch.arange(n_inp,device=x.device)) x = tok_emb + pos_emb x= self.drop_emb(x) x = self.tranf_blocks(x) x = self.final_norm(x) x = self.out_head(x) return x import torch.nn.functional as F def top_k_top_p_filtering(logits, top_k=0, top_p=0.9): if top_k > 0: values, _ = torch.topk(logits, top_k) min_values = values[:, -1].unsqueeze(-1) logits = torch.where( logits < min_values, torch.tensor(float("-inf"), device=logits.device), logits ) if top_p < 1.0: sorted_logits, sorted_indices = torch.sort( logits, descending=True ) cumulative_probs = torch.cumsum( F.softmax(sorted_logits, dim=-1), dim=-1 ) sorted_indices_to_remove = cumulative_probs > top_p sorted_indices_to_remove[:, 1:] = ( sorted_indices_to_remove[:, :-1].clone() ) sorted_indices_to_remove[:, 0] = False indices_to_remove = sorted_indices_to_remove.scatter( 1, sorted_indices, sorted_indices_to_remove ) logits = logits.masked_fill( indices_to_remove, float("-inf") ) return logits def apply_repetition_penalty(logits, generated_tokens, penalty=1.15): for token in set(generated_tokens.tolist()): logits[:, token] /= penalty return logits def generate_text( model, idx, max_new_tokens, context_size, temperature=0.65, top_k=30, top_p=0.9, repetition_penalty=1.15 ): model.eval() eos_token_id = 50256 with torch.no_grad(): for _ in range(max_new_tokens): idx_cond = idx[:, -context_size:] with torch.amp.autocast("cuda"): logits = model(idx_cond) logits = logits[:, -1, :] logits = apply_repetition_penalty( logits, idx[0], repetition_penalty ) logits = logits / temperature logits = top_k_top_p_filtering( logits, top_k=top_k, top_p=top_p ) probs = F.softmax(logits, dim=-1) idx_next = torch.multinomial( probs, num_samples=1 ) idx = torch.cat( (idx, idx_next), dim=1 ) if idx_next.item() == eos_token_id: break return idx def text_to_token_ids(text, tokenizer): encoded = tokenizer.encode(text) encoded_tensor = torch.tensor(encoded,device="cuda").unsqueeze(0) #1 return encoded_tensor def token_ids_to_text(token_ids, tokenizer): flat = token_ids.squeeze(0) return tokenizer.decode(flat.tolist()) def generate_and_print_sample( model, tokenizer, device, start_context ): context_size = model.pos_emb.weight.shape[0] encoded = text_to_token_ids( start_context, tokenizer ).to(device) token_ids = generate_text( model=model, idx=encoded, max_new_tokens=512, context_size=context_size, temperature=0.65, top_k=30, top_p=0.9, repetition_penalty=1.15 ) decoded_text = token_ids_to_text( token_ids, tokenizer ) print(decoded_text.replace("\n", " "))