| """Load and run the model in this folder: per question, a probability distribution over its options.
|
|
|
| from typed_decisions import TypedDecisions
|
| m = TypedDecisions.load(".") # this folder, or a downloaded snapshot of it
|
| m.predict("invoice_processing", state_dict_or_json)
|
| # {"discrepancy_severity": {"0": 0.71, "1": 0.22, ...}, "disposition": {...}, ...}
|
| """
|
| from __future__ import annotations
|
|
|
| import json
|
| from pathlib import Path
|
|
|
| import torch
|
| from safetensors.torch import load_file
|
| from torch import nn
|
| from transformers import AutoConfig, AutoModel, AutoTokenizer
|
|
|
|
|
| class TypedDecisions(nn.Module):
|
| def __init__(self, folder: Path, cfg: dict):
|
| super().__init__()
|
| enc_cfg = AutoConfig.from_pretrained(folder)
|
| self.enc = AutoModel.from_config(enc_cfg)
|
| h = enc_cfg.hidden_size
|
| self.schema: dict[str, dict[str, list[str]]] = cfg["schema"]
|
| self.heads = nn.ModuleDict({f"{wf}/{q}": nn.Linear(h, len(o)) for wf, qs in self.schema.items() for q, o in qs.items()})
|
| self.max_len = cfg["max_len"]
|
| self.tok = AutoTokenizer.from_pretrained(folder)
|
|
|
| @classmethod
|
| def load(cls, folder: str | Path, device: str = "cuda" if torch.cuda.is_available() else "cpu") -> "TypedDecisions":
|
| folder = Path(folder)
|
| m = cls(folder, json.loads((folder / "typed_decisions.json").read_text()))
|
| m.load_state_dict(load_file(str(folder / "model.safetensors")))
|
| return m.to(device).eval()
|
|
|
| @torch.no_grad()
|
| def predict(self, workflow: str, state: dict | str) -> dict[str, dict[str, float]]:
|
| """The input is the workflow name, a newline, then the state as compact JSON (as in training)."""
|
| s = json.dumps(json.loads(state) if isinstance(state, str) else state, ensure_ascii=False)
|
| dev = next(self.parameters()).device
|
| enc = self.tok([f"{workflow}\n{s}"], truncation=True, max_length=self.max_len, return_tensors="pt").to(dev)
|
| with torch.autocast(dev.type, dtype=torch.bfloat16, enabled=dev.type == "cuda"):
|
| hs = self.enc(**enc).last_hidden_state
|
| m = enc.attention_mask.unsqueeze(-1).to(hs.dtype)
|
| z = ((hs * m).sum(1) / m.sum(1)).float()
|
| return {q: dict(zip(opts, torch.softmax(self.heads[f"{workflow}/{q}"](z)[0], -1).tolist()))
|
| for q, opts in self.schema[workflow].items()}
|
|
|
