File size: 1,021 Bytes
e418394 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 | 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)
|