| """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)) | |