| |
| """Score text questions with this MLX checkpoint. |
| |
| The vision tower is stored at full precision in ``model.safetensors`` and is not |
| used here. Load requires ``mlx``, ``mlx-lm`` and ``transformers``. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import sys |
| from pathlib import Path |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parent)) |
|
|
| import d3_mlx |
|
|
|
|
| EXAMPLE = { |
| "state": "The order arrived damaged yesterday. The customer has a receipt and asks for a replacement today.", |
| "questions": { |
| "route": { |
| "type": "choice", |
| "instructions": "Which team should handle this request?", |
| "criteria": { |
| "returns": "Refunds, replacements and damaged deliveries", |
| "billing": "Payments, invoices and charges", |
| "technical": "Product setup and faults", |
| }, |
| }, |
| "receipt": {"type": "noul", "instructions": "Does the customer have a receipt?"}, |
| "urgency": { |
| "type": "score", |
| "instructions": "How urgent is this request?", |
| "criteria": ["Routine", "Soon", "Today"], |
| }, |
| }, |
| } |
|
|
|
|
| def option_labels(question: dict) -> list[str]: |
| rendered = d3_mlx.normalize_question(question) |
| if rendered["type"] == "noul": |
| return ["false", "true"] |
| return list(rendered["criteria"]) |
|
|
|
|
| def predict(model, readout, tokenizer, codes, state, questions: dict) -> dict: |
| answers = {} |
| for key, question in questions.items(): |
| text, n_options = d3_mlx.render_prompt(tokenizer, state, question, codes) |
| token_ids = tokenizer.encode(text, add_special_tokens=False) |
| probabilities = d3_mlx.probs_from_ids(model, readout, token_ids, n_options) |
| labels = option_labels(question) |
| pairs = {label: float(value) for label, value in zip(labels, probabilities)} |
| best = max(pairs, key=pairs.get) |
| answers[key] = {"best": best, "probabilities": pairs} |
| return answers |
|
|
|
|
| def main() -> None: |
| from transformers import AutoTokenizer |
|
|
| root = Path(__file__).resolve().parent |
| model, _config, decision, readout = d3_mlx.load_model(root) |
| tokenizer = AutoTokenizer.from_pretrained(str(root)) |
| payload = EXAMPLE |
| if len(sys.argv) > 1: |
| payload = json.loads(Path(sys.argv[1]).read_text(encoding="utf-8")) |
| answers = predict( |
| model, |
| readout, |
| tokenizer, |
| decision["codes"], |
| payload["state"], |
| payload["questions"], |
| ) |
| print(json.dumps(answers, ensure_ascii=False, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
