"""Run swahili-gpt-71m-instruct. Downloads everything from the Hub."""
from huggingface_hub import hf_hub_download
import sentencepiece as spm

REPO = "Benjamin-png/swahili-gpt-71m-instruct"
DEVICE = "cpu"   # or "cuda"

hf_hub_download(REPO, "modeling_kiswahili.py", local_dir=".")
from modeling_kiswahili import load_model, generate

weights = hf_hub_download(REPO, "pytorch_model.pt")
hf_hub_download(REPO, "model_config.json")
tok = hf_hub_download(REPO, "swahili_tokenizer.model")

model, cfg = load_model(weights, device=DEVICE)
sp = spm.SentencePieceProcessor(); sp.load(tok)


def ask(q, context=""):
    prompt = f"### Maagizo:\n{q}\n"
    if context:
        prompt += f"### Ingizo:\n{context}\n"
    prompt += "### Jibu:\n"
    out = generate(model, sp, prompt, device=DEVICE, max_new_tokens=120)
    return out.split("### Jibu:")[-1].split("### Maagizo:")[0].strip()


for q in [
    "Eleza umuhimu wa elimu kwa jamii.",
    "Andika sentensi fupi kuhusu teknolojia.",
    "Mji mkuu wa Kenya ni upi?",
]:
    print("SWALI :", q)
    print("JIBU  :", ask(q))
    print("-" * 60)