examples/quickstart.py

"""One context, multiple candidate answers; no text generation."""
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()

context = "CONTEXT:\nA customer disputes an invoice and asks for a billing correction."
options = [
    "route to billing support",
    "route to technical support",
    "route to sales",
]
decision = model.choose(tokenizer, context, options)
print(decision)  # choice, index, probabilities in supplied candidate order