"""Minimal inference example for TinyTurk v2.7d."""
import torch
import sentencepiece as spm
from huggingface_hub import hf_hub_download
from model import TinyTurkGPTV27, TinyTurkV27Config

REPO = "stunmuffin/TinyTurk-v2.7d"

ckpt_path = hf_hub_download(REPO, "pytorch_model.bin")
tok_path = hf_hub_download(REPO, "tokenizer_v11_bpe_512.model")

ckpt = torch.load(ckpt_path, map_location="cpu")
config = TinyTurkV27Config(**ckpt["config"])
model = TinyTurkGPTV27(config)
model.load_state_dict(ckpt["model_state_dict"])
model.eval()

sp = spm.SentencePieceProcessor(model_file=tok_path)

for prompt in ["Ali elmayi", "Bir sabah", "Kucuk kedi"]:
    ids = [sp.bos_id()] + sp.encode(prompt, out_type=int)
    input_ids = torch.tensor([ids])
    with torch.no_grad():
        out = model.generate(input_ids, max_new_tokens=40,
                             temperature=0.8, top_k=40)
    print("PROMPT:", prompt)
    print("OUTPUT:", sp.decode(out[0].tolist()))
    print("-" * 60)