"""Load and generate with Compactbot/tinystories-50m (custom GPT, 54.8M params).
Requires: torch, safetensors, tokenizers. GPU optional (CPU works, slower).
"""
import torch, torch.nn as nn, torch.nn.functional as F
from safetensors.torch import load_file
from tokenizers import Tokenizer

VOCAB,D,L,H,FFN,SEQ=8192,512,16,8,2048,512

class RMSNorm(nn.Module):
    def __init__(self,d): super().__init__(); self.w=nn.Parameter(torch.ones(d))
    def forward(self,x): return self.w*x*torch.rsqrt(x.float().pow(2).mean(-1,keepdim=True)+1e-6)

class Block(nn.Module):
    def __init__(self,d,h):
        super().__init__()
        self.ln1=RMSNorm(d); self.ln2=RMSNorm(d)
        self.qkv=nn.Linear(d,3*d,bias=False); self.proj=nn.Linear(d,d,bias=False)
        self.fc1=nn.Linear(d,FFN,bias=False); self.fc2=nn.Linear(FFN,d,bias=False)
        self.h,self.d=h,d
    def forward(self,x):
        B,T,Dd=x.shape; h=self.ln1(x)
        qkv=self.qkv(h).view(B,T,3,self.h,Dd//self.h).transpose(2,1)
        q,k,v=qkv[:,0].transpose(1,2),qkv[:,1].transpose(1,2),qkv[:,2].transpose(1,2)
        att=F.scaled_dot_product_attention(q,k,v,is_causal=True)
        x=x+self.proj(att.transpose(1,2).reshape(B,T,Dd))
        return x+self.fc2(F.gelu(self.fc1(self.ln2(x))))

class GPT(nn.Module):
    def __init__(self):
        super().__init__()
        self.tok=nn.Embedding(VOCAB,D); self.pos=nn.Embedding(SEQ,D)
        self.blocks=nn.ModuleList([Block(D,H) for _ in range(L)]); self.ln_f=RMSNorm(D)
    def forward(self,x,targets=None):
        h=self.tok(x)+self.pos(torch.arange(x.shape[1],device=x.device))
        for b in self.blocks: h=b(h)
        logits=self.ln_f(h)@self.tok.weight.t()
        if targets is not None: return F.cross_entropy(logits.view(-1,VOCAB),targets.view(-1))
        return logits

def load(repo_dir=".",device="cuda" if torch.cuda.is_available() else "cpu"):
    m=GPT().to(device); m.load_state_dict(load_file(f"{repo_dir}/model.safetensors"),strict=True)
    m.eval()
    return m,Tokenizer.from_file(f"{repo_dir}/tokenizer.json")

def generate(model,tok,prompt,max_new=120,seed=0,temperature=0.8,device="cuda"):
    g=torch.Generator(device=device).manual_seed(seed)
    ids=torch.tensor([tok.encode(prompt).ids],device=device)
    if ids.shape[1]>SEQ-4: ids=ids[:,-(SEQ-4):]
    with torch.no_grad():
        for _ in range(max_new):
            with torch.autocast(device_type=device,dtype=torch.bfloat16):
                logits=model(ids)
            p=torch.softmax(logits[:,-1].float()/temperature,dim=-1)
            nxt=torch.multinomial(p,1,generator=g); ids=torch.cat([ids,nxt],1)
            if nxt.item()==1: break
    return tok.decode(ids[0].tolist())

if __name__=="__main__":
    m,t=load(".")
    print("params:",sum(p.numel() for p in m.parameters()))
    print(generate(m,t,"Once upon a time, there was a little girl named Lily.",seed=0))