"""Run Pina 500M. Needs: pip install torch safetensors tiktoken huggingface_hub"""
import json, pickle, torch
from safetensors.torch import load_model
from huggingface_hub import hf_hub_download
from modeling_pina import SparseLM, SparseLMConfig

REPO, TOK = "procmarco/pina-500m", "procmarco/ita-en-code-bpe-48k"
cfg = SparseLMConfig(**{k: v for k, v in json.load(open(hf_hub_download(REPO, "config.json"))).items() if k != "architecture"})
model = SparseLM(cfg).eval()
if torch.cuda.is_available():
    model = model.cuda().to(torch.bfloat16)
load_model(model, hf_hub_download(REPO, "model.safetensors"))

enc = pickle.load(open(hf_hub_download(TOK, "tokenizer.pkl"), "rb"))
bos, eos = enc.encode_single_token("<|bos|>"), enc.encode_single_token("<|eos|>")

for prompt in ["L'Italia è un paese famoso per", "The theory of relativity states that", "La ricetta della carbonara:"]:
    ids = [bos] + enc.encode_ordinary(prompt)
    out = model.generate(ids, max_new_tokens=60, temperature=0.8, top_k=40, eos_id=eos)
    print(f"\n>>> {prompt}\n{enc.decode(out[1:])}")