| """Download the Hugging Face checkpoint and run one typed Choice prediction.""" |
|
|
| from __future__ import annotations |
|
|
| import sys |
| from pathlib import Path |
|
|
| import torch |
| from huggingface_hub import snapshot_download |
|
|
|
|
| REPO_ID = "Syigen/rgcm" |
|
|
|
|
| def main() -> None: |
| repo_path = Path(snapshot_download(REPO_ID)) |
| |
| sys.path.insert(0, str(repo_path)) |
|
|
| from concept_memory import ConceptMemoryModel |
| from concept_memory.types import Question |
|
|
| checkpoint = repo_path / "model.pt" |
| payload = torch.load(checkpoint, map_location="cpu", weights_only=False) |
|
|
| model = ConceptMemoryModel.from_pretrained( |
| payload["encoder"], |
| decision_size=payload["decision_size"], |
| max_length=payload["max_length"], |
| ) |
| model.load_state_dict(payload["model_state_dict"]) |
| model.eval() |
|
|
| question = Question( |
| id="intent", |
| type="choice", |
| instructions="Which team should handle this customer request?", |
| criteria=[ |
| {"id": "billing", "description": "Payment or subscription issues"}, |
| {"id": "technical", "description": "Bugs or integration problems"}, |
| {"id": "sales", "description": "Pricing or account questions"}, |
| ], |
| ) |
|
|
| with torch.inference_mode(): |
| result = model.predict( |
| "Stripe integration keeps failing after several attempts.", [question] |
| )[0] |
|
|
| criteria = question.normalized_criteria() |
| predicted = criteria[result.prediction.item()] |
| print(f"Prediction: {predicted.id} - {predicted.description}") |
| print("Probabilities:", result.probabilities.tolist()) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
