| """Edit STATE, QUESTION, and OPTIONS, then run: python run.py.""" |
| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
|
|
| import torch |
| from safetensors.torch import load_file |
| from transformers import AutoTokenizer |
|
|
| from data_utils import make_collate |
| from decision_model import DecisionModel, ModelConfig |
|
|
| |
| STATE = "I was charged twice for the same purchase." |
| QUESTION = "Which issue is this?" |
| OPTIONS = ["delivery problem", "duplicate charge", "wrong item"] |
| |
|
|
|
|
| def load_decision_maker(model_dir: str | Path): |
| model_dir = Path(model_dir) |
| config = json.loads((model_dir / "model_config.json").read_text(encoding="utf-8")) |
| tokenizer = AutoTokenizer.from_pretrained(model_dir / "tokenizer") |
| model = DecisionModel(ModelConfig(**config["model_config"])) |
| model.resize_token_embeddings(len(tokenizer)) |
| model.load_state_dict(load_file(model_dir / "model.safetensors")) |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| return model.to(device).eval(), tokenizer, config, device |
|
|
|
|
| def decide(model, tokenizer, config, device, state: str, question: str, options: list[str]): |
| row = {"id": "request", "state": state, "question": question, "options": options, "label": 0} |
| batch = make_collate(tokenizer, config["max_length"])([row]) |
| tensors = {key: value.to(device) for key, value in batch.items() if isinstance(value, torch.Tensor) and key != "labels"} |
| temperature = float(config.get("temperature", 1.0)) |
| with torch.no_grad(): |
| probabilities = (model(**tensors)[0] / temperature).softmax(-1).cpu() |
| distribution = {option: round(probabilities[index].item(), 6) for index, option in enumerate(options)} |
| return {"choice": options[int(probabilities.argmax())], "temperature": temperature, "probabilities": distribution} |
|
|
|
|
| if __name__ == "__main__": |
| root = Path(__file__).parent |
| model, tokenizer, config, device = load_decision_maker(root) |
| print(json.dumps(decide(model, tokenizer, config, device, STATE, QUESTION, OPTIONS), indent=2)) |
|
|
