"""Self-contained PyTorch implementation of tiny-edu-114m (plain PyTorch, no transformers modeling code needed).

    import torch
    from safetensors.torch import load_file
    from modeling_tiny_edu import TinyEdu, generate
    model = TinyEdu(); model.load_state_dict(load_file("model.safetensors"), strict=False); model.eval()   # strict=False: the output head is tied to the embedding
    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained("gpt2")
    print(generate(model, tok, "The capital of France is", max_new=30))
"""
import torch, torch.nn as nn, torch.nn.functional as F

D, NH, NKV, NL, V, FF = 768, 12, 4, 12, 50257, 3072       # width, query heads, key/value heads, layers, vocabulary, MLP width
HD = D // NH                                               # head dimension 64


class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.ln_1, self.ln_2 = nn.LayerNorm(D, eps=1e-5), nn.LayerNorm(D, eps=1e-5)
        self.attn = nn.Module(); self.attn.q_proj = nn.Linear(D, NH * HD, bias=False); self.attn.k_proj = nn.Linear(D, NKV * HD, bias=False)
        self.attn.v_proj = nn.Linear(D, NKV * HD, bias=False); self.attn.out_proj = nn.Linear(NH * HD, D, bias=False)
        self.mlp = nn.Module(); self.mlp.c_fc = nn.Linear(D, FF, bias=False); self.mlp.c_proj = nn.Linear(FF, D, bias=False)

    def forward(self, x, cos, sin):
        B, T, _ = x.shape
        h = self.ln_1(x)
        q = self.attn.q_proj(h).view(B, T, NH, HD); k = self.attn.k_proj(h).view(B, T, NKV, HD); v = self.attn.v_proj(h).view(B, T, NKV, HD)

        def rot(t):                                         # rotary embedding, half-split rotation
            t1, t2 = t[..., :HD // 2], t[..., HD // 2:]
            c, s = cos[None, :T, None, :].to(t.dtype), sin[None, :T, None, :].to(t.dtype)
            return torch.cat([t1 * c - t2 * s, t2 * c + t1 * s], -1)
        q, k = rot(q), rot(k)
        k = k.repeat_interleave(NH // NKV, 2); v = v.repeat_interleave(NH // NKV, 2)
        o = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), is_causal=True).transpose(1, 2).reshape(B, T, NH * HD)
        x = x + self.attn.out_proj(o)
        return x + self.mlp.c_proj(F.gelu(self.mlp.c_fc(self.ln_2(x))))


class TinyEdu(nn.Module):
    def __init__(self):
        super().__init__()
        self.transformer = nn.Module(); self.transformer.wte = nn.Embedding(V, D); self.transformer.h = nn.ModuleList([Block() for _ in range(NL)]); self.transformer.ln_f = nn.LayerNorm(D, eps=1e-5)
        inv = 1.0 / (10000 ** (torch.arange(0, HD, 2).float() / HD)); f = torch.outer(torch.arange(4096).float(), inv)
        self.register_buffer("cos", f.cos(), persistent=False); self.register_buffer("sin", f.sin(), persistent=False)

    def forward(self, ids):                                 # (B, T) token ids -> (B, T, 50257) logits; the output head shares the embedding matrix
        x = self.transformer.wte(ids)
        for b in self.transformer.h:
            x = b(x, self.cos, self.sin)
        return self.transformer.ln_f(x) @ self.transformer.wte.weight.T


@torch.no_grad()
def generate(model, tok, prompt, max_new=60, temperature=0.8, top_k=50, device=None):
    """Plain sampling without a cache (fine for a model this small); temperature 0 = greedy."""
    device = device or next(model.parameters()).device
    ids = torch.tensor([tok.encode(prompt)], device=device)
    for _ in range(max_new):
        lg = model(ids[:, -1024:])[:, -1].float()
        if temperature == 0:
            nxt = lg.argmax(-1, keepdim=True)
        else:
            lg = lg / temperature
            if top_k:
                lg[lg < lg.topk(top_k).values[:, -1:]] = -float("inf")
            nxt = torch.multinomial(lg.softmax(-1), 1)
        ids = torch.cat([ids, nxt], 1)
        if int(nxt) == 50256:
            break
    return tok.decode(ids[0].tolist())