"""
Minimal example: load Quartz Micro Preview Base's weights and generate text.

This is a custom PyTorch architecture (DeepSeek-style fine-grained MoE, see
model.py / config.py in this repo) -- not a registered transformers
architecture, so there's no AutoModel support. Load it directly:

    pip install torch safetensors tokenizers
    python load_model.py

The tokenizer is a from-scratch 32K-vocab byte-level BPE (tokenizer/
vocab.json + merges.txt) trained on this model's own pretraining corpus --
it is NOT compatible with any other model's tokenizer, and no other
tokenizer is compatible with these weights.
"""
import os

import torch
from safetensors.torch import load_file
from tokenizers import ByteLevelBPETokenizer

from config import MoEConfig
from model import MoELanguageModel

HERE = os.path.dirname(__file__)


def load():
    tok = ByteLevelBPETokenizer(
        os.path.join(HERE, "tokenizer", "vocab.json"),
        os.path.join(HERE, "tokenizer", "merges.txt"),
    )

    cfg = MoEConfig()
    model = MoELanguageModel(cfg)
    state_dict = load_file(os.path.join(HERE, "model.safetensors"))
    model.load_state_dict(state_dict)
    model.eval()
    return model, tok


@torch.no_grad()
def generate(model, tok, prompt: str, max_new_tokens: int = 80, temperature: float = 0.8):
    BOS_ID, EOS_ID = 1, 2
    ids = [BOS_ID] + tok.encode(prompt).ids
    x = torch.tensor([ids], dtype=torch.long)
    for _ in range(max_new_tokens):
        logits, _ = model(x)
        next_logits = logits[0, -1] / temperature
        probs = torch.softmax(next_logits, dim=-1)
        next_id = torch.multinomial(probs, 1).item()
        if next_id == EOS_ID:
            break
        x = torch.cat([x, torch.tensor([[next_id]])], dim=1)
    return tok.decode(x[0, 1:].tolist())


if __name__ == "__main__":
    model, tok = load()
    print(generate(model, tok, "The history of the Roman Empire"))