| """DecisionAgent: `predict(state, questions)` -> Laya/Jev-shaped answers.
|
|
|
| Output (same schema as Laya, which is schema-identical to TypeSafe Jev):
|
| {"model": name,
|
| "answers": {qid: {"type": "choice", "choice": ..., "probabilities": {...},
|
| "confidence": ..., "answer_confidence": ..., "action": {...}}, ...},
|
| "usage": {"input_tokens": N, "output_tokens": 0}}
|
|
|
| All questions for all states are scored in one batched forward pass per chunk.
|
| """
|
| import os
|
| from typing import Any, Dict, List, Optional, Union
|
|
|
| import numpy as np
|
| import torch
|
|
|
| from decision_core import (QTYPES, answer_confidence, assemble, check_question, clamp_temperature,
|
| collate, confidence_from_probs, load_decision_model, load_tokenizer,
|
| temp_bucket, to_internal, tokenize_question, tokenize_state)
|
|
|
|
|
| class DecisionAgent:
|
| def __init__(self, model_path: str, code_dir: Optional[str] = None, device: str = "cpu",
|
| max_rows: int = 16):
|
| code_dir = code_dir or os.path.dirname(os.path.abspath(__file__))
|
| self.device = torch.device(device)
|
| self.model, self.cfg, _ = load_decision_model(code_dir, model_path, self.device)
|
| self.tok = load_tokenizer(code_dir)
|
| self.name = self.cfg.get("model_name", "byrne-decisions")
|
| self.max_len = self.cfg.get("max_len", 1024)
|
| self.head_max_len = self.cfg.get("head_max_len", 256)
|
| self.temperature = [clamp_temperature(t) for t in self.cfg.get("temperature", [1.0, 1.0, 1.0])]
|
| self.temperature_by_options = {k: clamp_temperature(v)
|
| for k, v in self.cfg.get("temperature_by_options", {}).items()}
|
| self.max_rows = max_rows
|
| self.amp = self.device.type == "cuda"
|
|
|
| @torch.no_grad()
|
| def _logits(self, rows: List[Dict]):
|
| out_l, out_a = [], []
|
| for i in range(0, len(rows), self.max_rows):
|
| b = collate(rows[i:i + self.max_rows], self.tok.pad_token_id)
|
| with torch.autocast(self.device.type, dtype=torch.bfloat16, enabled=self.amp):
|
| logits, act = self.model(*(b[k].to(self.device) for k in
|
| ("input_ids", "attention_mask", "marker_pos", "marker_mask", "qtype")))
|
| out_l.extend(logits.float().cpu().numpy())
|
| out_a.extend(torch.softmax(act.float(), -1).cpu().numpy())
|
| return out_l, out_a
|
|
|
| def _decode(self, q: Dict, z_row, act_row, k: int) -> Dict[str, Any]:
|
| qt = QTYPES[q["t"]]
|
| t_scale = self.temperature_by_options.get(temp_bucket(qt, k), self.temperature[qt])
|
| z = z_row[:k] / t_scale
|
| p = np.exp(z - z.max())
|
| p = p / p.sum()
|
| ans_conf = round(answer_confidence(p, k), 4)
|
| ext = {"act_probability": round(float(act_row[0]), 4)}
|
| if q["t"] == "choice":
|
| keys = list(q["crit"].keys())
|
| return {"type": "choice", "choice": keys[int(p.argmax())],
|
| "probabilities": {kk: round(float(v), 4) for kk, v in zip(keys, p)},
|
| "confidence": round(confidence_from_probs(p, k), 4),
|
| "answer_confidence": ans_conf, "action": ext}
|
| if q["t"] == "score":
|
| return {"type": "score", "score": round(float((np.arange(k) * p).sum()), 4),
|
| "legend": {str(i): c for i, c in enumerate(q["crit"])},
|
| "probabilities": {str(i): round(float(v), 4) for i, v in enumerate(p)},
|
| "confidence": round(confidence_from_probs(p, k), 4),
|
| "answer_confidence": ans_conf, "action": ext}
|
| return {"type": "noul", "noul": round(float(p[1]), 4),
|
| "confidence": round(max(float(p[1]), 1.0 - float(p[1])), 4),
|
| "answer_confidence": ans_conf, "action": ext}
|
|
|
| def predict_batch(self, states: List[Union[str, dict, list]], questions: Dict[str, Dict]) -> List[Dict]:
|
| if not isinstance(questions, dict):
|
| raise ValueError("questions must be a dict of question_id -> definition")
|
| for qid, qdef in questions.items():
|
| check_question(qid, qdef)
|
| ids = list(questions)
|
| internal = {qid: to_internal(questions[qid]) for qid in ids}
|
| if not ids:
|
| return [{"model": self.name, "answers": {}, "usage": {"input_tokens": 0, "output_tokens": 0}}
|
| for _ in states]
|
| qtoks = {qid: tokenize_question(self.tok, internal[qid]) for qid in ids}
|
| rows, per_state = [], []
|
| for st in states:
|
| st_ids = tokenize_state(self.tok, st)
|
| n_in = 0
|
| for qid in ids:
|
| seq, markers = assemble(self.tok.bos_token_id, st_ids, qtoks[qid], self.max_len,
|
| self.head_max_len, truncate_left=isinstance(st, list))
|
| if len(markers) != len(qtoks[qid]["opts"]):
|
| raise ValueError("question %r: options exceed head_max_len=%d" % (qid, self.head_max_len))
|
| rows.append({"ids": seq, "markers": markers, "qtype": QTYPES[internal[qid]["t"]]})
|
| n_in += len(seq)
|
| per_state.append(n_in)
|
| logits, act = self._logits(rows)
|
| results, r = [], 0
|
| for s_i in range(len(states)):
|
| answers = {}
|
| for qid in ids:
|
| answers[qid] = self._decode(internal[qid], logits[r], act[r], len(rows[r]["markers"]))
|
| r += 1
|
| results.append({"model": self.name, "answers": answers,
|
| "usage": {"input_tokens": per_state[s_i], "output_tokens": 0},
|
|
|
|
|
| "rounding": {"probabilityDecimals": 4}})
|
| return results
|
|
|
| def predict(self, state, questions: Dict[str, Dict]) -> Dict:
|
| return self.predict_batch([state], questions)[0]
|
|
|
| system_one = predict
|
|
|
