"""Small character-level GPT-like Transformer; initialized from random weights."""
import torch
from torch import nn

def encode(text, vocab): return [vocab[c] for c in text if c in vocab]
def decode(ids, vocab):
    inverse = {v:k for k,v in vocab.items()}
    return ''.join(inverse.get(int(i), '') for i in ids)

class Head(nn.Module):
    def __init__(self, n_embd, head_size, block_size):
        super().__init__(); self.key=nn.Linear(n_embd,head_size,bias=False); self.query=nn.Linear(n_embd,head_size,bias=False); self.value=nn.Linear(n_embd,head_size,bias=False)
        self.register_buffer('tril',torch.tril(torch.ones(block_size,block_size)))
    def forward(self,x):
        B,T,C=x.shape; k=self.key(x); q=self.query(x); weights=q@k.transpose(-2,-1)*(k.shape[-1]**-0.5); weights=weights.masked_fill(self.tril[:T,:T]==0,float('-inf')); weights=torch.softmax(weights,dim=-1); return weights@self.value(x)
class MultiHead(nn.Module):
    def __init__(self,n_embd,n_head,block_size):
        super().__init__(); size=n_embd//n_head; self.heads=nn.ModuleList([Head(n_embd,size,block_size) for _ in range(n_head)]); self.proj=nn.Linear(n_embd,n_embd)
    def forward(self,x): return self.proj(torch.cat([h(x) for h in self.heads],dim=-1))
class Block(nn.Module):
    def __init__(self,n_embd,n_head,block_size):
        super().__init__(); self.sa=MultiHead(n_embd,n_head,block_size); self.ff=nn.Sequential(nn.Linear(n_embd,4*n_embd),nn.ReLU(),nn.Linear(4*n_embd,n_embd)); self.ln1=nn.LayerNorm(n_embd); self.ln2=nn.LayerNorm(n_embd)
    def forward(self,x): x=x+self.sa(self.ln1(x)); return x+self.ff(self.ln2(x))
class TinyGPT(nn.Module):
    def __init__(self,vocab_size,block_size=128,n_embd=128,n_head=4,n_layer=3):
        super().__init__(); self.block_size=block_size; self.token=nn.Embedding(vocab_size,n_embd); self.position=nn.Embedding(block_size,n_embd); self.blocks=nn.Sequential(*[Block(n_embd,n_head,block_size) for _ in range(n_layer)]); self.ln=nn.LayerNorm(n_embd); self.lm_head=nn.Linear(n_embd,vocab_size)
    def forward(self,idx):
        B,T=idx.shape; x=self.token(idx)+self.position(torch.arange(T,device=idx.device)); return self.lm_head(self.ln(self.blocks(x)))