"""Load morena-0.5b-mini-sft and sample from it.

Morena is not a transformers architecture; it is `Transformer` from modeling_morena.py, and that
class has NO .generate() and a training-shaped forward:

    forward(idx, positions, cu_seqlens, max_seqlen, mask)

For single-sequence eval pass cu_seqlens=None and mask=None, which is plain causal SDPA. The
sampler below is the one the eval harness uses (no KV cache; fine for a couple hundred tokens).

    pip install torch safetensors tokenizers
    python load_example.py
"""
import json, torch
from safetensors.torch import load_file
from tokenizers import Tokenizer
import modeling_morena as M

cfg = json.load(open("config.json"))
model = M.Transformer(M.ModelConfig(**cfg["model"]), "sdpa")     # "sdpa" = plain causal attention
model.load_state_dict(load_file("model.safetensors"), strict=True)
model = model.to(torch.bfloat16).cuda().eval()
tok = Tokenizer.from_file("tokenizer.json")

EOS = tok.token_to_id("<eos>")
# ---------------------------------------------------------------------------
# CHAT MARKERS -- DO NOT "MODERNISE" OR "CORRECT" THESE.
#
# morena-0.5b-mini-sft was SFT'd with the sft5/sft10-era chat format: single reserved tokens
# <reserved_0> (user, id 3) and <reserved_1> (assistant, id 4). They are real single tokens,
# not the literal angle-bracket strings they look like.
# ---------------------------------------------------------------------------
USER, ASSISTANT = "<reserved_0>", "<reserved_1>"


@torch.no_grad()
def generate(prompt, max_new_tokens=100, temperature=0.7, top_p=0.9):
    ids = [EOS] + tok.encode(prompt, add_special_tokens=False).ids   # docs start with <eos>
    x = torch.tensor([ids], device="cuda")
    out = []
    for _ in range(max_new_tokens):
        pos = torch.arange(x.shape[1], device=x.device).unsqueeze(0)
        logits = model(x, pos, None, x.shape[1], None)[:, -1, :].float()
        probs = torch.softmax(logits / max(temperature, 1e-5), dim=-1)
        sp, si = torch.sort(probs, descending=True, dim=-1)
        sp[sp.cumsum(-1) - sp > top_p] = 0.0
        nxt = si.gather(-1, torch.multinomial(sp / sp.sum(-1, keepdim=True), 1))
        t = int(nxt[0, 0])
        if t == EOS:
            break
        out.append(t)
        x = torch.cat([x, nxt], dim=1)
    return tok.decode(out, skip_special_tokens=True)


print(generate(f"{USER}\nNdeipi guta guru reZimbabwe?\n{ASSISTANT}\n"))