File size: 2,194 Bytes
c800ccd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | """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
# ----- Edit only these values -------------------------------------------------
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))
|