"""Reuse candidate embeddings across multiple calls with the same options."""
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()

options = [
    "route to billing support",
    "route to technical support",
    "route to sales",
]
candidate_vectors = model.encode_candidates(tokenizer, options)

for context in (
    "A customer disputes an invoice amount.",
    "A customer reports a failed software update.",
):
    context_vector = model.encode_context(tokenizer, context)
    print(model.choose_cached(context_vector, candidate_vectors, options))