"""Structured decisions keep candidate IDs out of the option encoder."""
import torch
from transformers import AutoModel, AutoTokenizer

MODEL = "Q1z/Pivot"
tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
model = AutoModel.from_pretrained(MODEL, trust_remote_code=True, dtype=torch.float32).eval()

decision = model.decide_native(
    tokenizer,
    "A customer disputes an invoice and asks for a billing correction.",
    [
        {"id": "billing", "text": "route to billing support"},
        {"id": "technical", "text": "route to technical support"},
        {"id": "sales", "text": "route to sales"},
        {"id": "abstain", "text": "none of these routes is appropriate", "kind": "abstain"},
    ],
)
print(decision["selected"])
print(decision["prob_vector"])