import torch
from ouro_lens import from_ouro, load_ouro
torch.set_grad_enabled(False)
hf, tok = load_ouro()
hf = hf.cuda()
print("model loaded", flush=True)

def gen(text, n=48):
    ids = tok(text, return_tensors="pt").input_ids.cuda()
    out = []
    for _ in range(n):
        lg = hf(ids, use_cache=False).logits[0, -1]
        t = int(lg.argmax())
        out.append(t)
        if t == tok.eos_token_id:
            break
        ids = torch.cat([ids, torch.tensor([[t]]).cuda()], dim=1)
    return tok.decode(out)

raw = "Fact: The currency used in the country shaped like a boot is"
print("RAW:", repr(gen(raw)), flush=True)
msgs = [{"role":"user","content":"What is the capital of France? Answer in one word."}]
print("CHAT(no think):", repr(gen(tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True))), flush=True)
print("CHAT(think):", repr(gen(tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True, enable_thinking=True))), flush=True)
print("DONE_GENCHECK", flush=True)