"""Ask Jegue questions with the Hugging Face class (one forward pass, no text generated).

Needs one CUDA GPU (about 20 GB), torch, transformers >= 5, peft. From the repository root:

    python examples/ask_hf.py [direct|routed]
"""
import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "code"))
from jegue import Jegue  # noqa: E402

model = Jegue(mode=sys.argv[1] if len(sys.argv) > 1 else "direct")

state = "Email from a customer: 'My parcel was due on Monday and it is Friday. Tracking has not changed since Tuesday. I want my money back.'"
questions = {
    "intent": {
        "type": "choice",
        "instructions": "What does the customer mainly want?",
        "criteria": {"refund": "A refund", "status": "An update on the delivery", "complaint": "To complain about the service", "other": "Something else"},
    },
    "urgent": {
        "type": "noul",   # yes/no: returns the probability of yes
        "instructions": "Should this email be escalated to a human agent today?",
    },
}

response = model.ask(state, questions)
for qid, answer in response["answers"].items():
    if answer["type"] == "choice":
        ranked = sorted(answer["probabilities"].items(), key=lambda kv: -kv[1])
        print(f"{qid}: {answer['choice']}  " + "  ".join(f"{k}={v:.2f}" for k, v in ranked))
    else:
        print(f"{qid}: P(yes)={answer['noul']:.2f}")