"""Matrix-BIOS-Italo-0.1 — compact (41.5M) Italian text-generation preview.
Custom architecture (loads via trust_remote_code) with a word-level vocabulary.
NOTE: this is a v0.1 research preview — it demonstrates the on-prem pipeline and
footprint, not production fluency.  pip install torch transformers huggingface_hub
"""
import json, torch
from huggingface_hub import hf_hub_download
from transformers import AutoModelForCausalLM

REPO = "ruslanmv/Matrix-BIOS-Italo-0.1"
model = AutoModelForCausalLM.from_pretrained(REPO, trust_remote_code=True).eval()

# Italo ships a word-level vocab.json (no standard tokenizer file).
vocab = json.load(open(hf_hub_download(REPO, "vocab.json")))
inv = {i: w for w, i in vocab.items()}
UNK, PAD = vocab.get("<unk>", 0), vocab.get("<pad>", 1)
encode = lambda t: [vocab.get(w, UNK) for w in t.lower().split()]
decode = lambda ids: " ".join(inv.get(int(i), "<unk>") for i in ids)

def generate(prompt: str, n: int = 12) -> str:
    ids = torch.tensor([encode(prompt)])
    with torch.no_grad():
        out = model.generate(ids, max_new_tokens=n, do_sample=False, pad_token_id=PAD)
    return decode(out[0])

if __name__ == "__main__":
    for p in ["la capitale d' italia", "matrix bios e"]:
        print(f"[{p!r}] -> {generate(p)}")