File size: 751 Bytes
7c85c7e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 | """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))
|