| """Inference for an EmbeddingGemma 2 decision model (Clef-style head). |
| |
| Usage: |
| from decision_model import DecisionModel |
| dm = DecisionModel("username/embeddinggemma2-decision") # or a local folder |
| dm.decide("Hi, I was charged twice...", { |
| "team": {"type": "choice", "instructions": "Which team should handle this?", |
| "criteria": {"billing": "payments, refunds", "technical": "bugs", "sales": "pricing"}}, |
| "refund": {"type": "noul", "instructions": "The customer asks for a refund."}, |
| "urgency": {"type": "score", "instructions": "How urgent is this?", |
| "criteria": ["Can wait", "Normal queue", "This week", "Today"]}, |
| }) |
| |
| Run in float32 or bfloat16 only. EmbeddingGemma 2 returns NaN in float16. |
| """ |
| import os |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| from transformers import AutoModel, AutoTokenizer |
|
|
|
|
| class DecisionHead(nn.Module): |
| """Score of an option = MLP(option vector, question vector).""" |
|
|
| def __init__(self, d, hidden=512): |
| super().__init__() |
| self.norm = nn.LayerNorm(d) |
| self.w_opt, self.w_q = nn.Linear(d, hidden), nn.Linear(d, hidden) |
| self.w_int = nn.Linear(d, hidden, bias=False) |
| self.out = nn.Sequential(nn.GELU(), nn.Dropout(0.1), nn.Linear(hidden, 1)) |
|
|
| def forward(self, opt_vecs, q_vec): |
| o, q = self.norm(opt_vecs), self.norm(q_vec).unsqueeze(0) |
| z = self.w_opt(o) + self.w_q(q) + self.w_int(o * q) |
| return self.out(z).squeeze(-1) |
|
|
|
|
| def option_list(spec): |
| """choice: criteria {key: description}; score: criteria [level descriptions]; noul: false/true.""" |
| t, crit = spec.get("type", "choice"), spec.get("criteria") |
| if isinstance(crit, dict) and crit: |
| keys = list(crit.keys()) |
| return keys, [str(crit[k]) for k in keys] |
| if isinstance(crit, list) and crit: |
| return [str(i) for i in range(len(crit))], [str(c) for c in crit] |
| if t == "noul": |
| return ["false", "true"], ["No.", "Yes."] |
| if t == "score": |
| n = int(spec.get("levels", 4)) |
| return [str(i) for i in range(n)], [f"Level {i}" for i in range(n)] |
| raise ValueError(f"Cannot determine options for question: {spec}") |
|
|
|
|
| def question_text(qtype, instructions): |
| if qtype == "noul": |
| return "\nstatement (true or false): " + instructions |
| if qtype == "score": |
| return "\nrating question: " + instructions |
| return "\nquestion: " + instructions |
|
|
|
|
| def option_text(key, desc): |
| return f"{key}: {desc}" if desc and desc != key else str(key) |
|
|
|
|
| def temp_for(temps, domain, qtype): |
| for g in (domain, "typed", "extra"): |
| if f"{g}|{qtype}" in temps: |
| return temps[f"{g}|{qtype}"] |
| return 1.0 |
|
|
|
|
| def softmax_np(x, T=1.0): |
| z = x / T |
| z = z - z.max() |
| e = np.exp(z) |
| return e / e.sum() |
|
|
|
|
| class DecisionModel: |
| def __init__(self, path_or_repo, device=None, token=None): |
| path = path_or_repo |
| if not os.path.isdir(path): |
| from huggingface_hub import snapshot_download |
| path = snapshot_download(path_or_repo, token=token) |
| ck = torch.load(os.path.join(path, "head.pt"), map_location="cpu") |
| cfg = ck.get("config", {}) |
| self.prefix, self.max_len, self.temps = ck["prefix"], int(cfg.get("max_len", 1024)), ck.get("temps", {}) |
| self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") |
| self.bf16 = self.device == "cuda" and torch.cuda.is_bf16_supported() |
| self.tok = AutoTokenizer.from_pretrained(os.path.join(path, "encoder")) |
| self.enc = AutoModel.from_pretrained(os.path.join(path, "encoder"), torch_dtype=torch.float32) |
| self.enc.to(self.device).eval() |
| self.head = DecisionHead(ck["d_model"]) |
| self.head.load_state_dict(ck["head"]) |
| self.head.to(self.device).eval() |
| self.lead, self.trail = self._special_wrap() |
| self.pad_id = self.tok.pad_token_id if self.tok.pad_token_id is not None else 0 |
| self.mark = self._ids("\n- ") |
|
|
| def _ids(self, s): |
| return self.tok(s, add_special_tokens=False)["input_ids"] |
|
|
| def _special_wrap(self): |
| full = self.tok("hello world")["input_ids"] |
| core = self._ids("hello world") |
| for p in range(len(full) - len(core) + 1): |
| if full[p:p + len(core)] == core: |
| return full[:p], full[p + len(core):] |
| return [], [] |
|
|
| def _build(self, text, qs): |
| parts = [] |
| for q in qs: |
| qids = self._ids(question_text(q["type"], q["instructions"])) |
| opts = [self._ids(option_text(k, d)) for k, d in zip(q["keys"], q["descs"])] |
| parts.append((qids, opts)) |
| q_total = sum(len(qi) + sum(len(self.mark) + len(o) for o in os_) for qi, os_ in parts) |
| prefix = self._ids(self.prefix) |
| budget = self.max_len - len(self.lead) - len(self.trail) - len(prefix) - q_total |
| if budget < 16: |
| raise ValueError("Questions are too long for max_len") |
| ids = list(self.lead) + prefix + self._ids(text)[:budget] |
| q_spans, o_spans = [], [] |
| for qids, opts in parts: |
| s = len(ids) |
| ids += qids |
| q_spans.append((s, len(ids))) |
| spans = [] |
| for o in opts: |
| ids += self.mark |
| s = len(ids) |
| ids += o |
| spans.append((s, len(ids))) |
| o_spans.append(spans) |
| return ids + list(self.trail), q_spans, o_spans |
|
|
| @torch.no_grad() |
| def decide(self, text, questions, domain="extra"): |
| """domain: "extra" for natural text (messages, emails), "typed" for JSON state like typed-decisions.""" |
| qs = [] |
| for name, spec in questions.items(): |
| keys, descs = option_list(spec) |
| qs.append(dict(name=name, type=spec.get("type", "choice"), |
| instructions=spec.get("instructions", ""), keys=keys, descs=descs)) |
| ids, q_spans, o_spans = self._build(text if isinstance(text, str) else str(text), qs) |
| x = torch.tensor([ids], device=self.device) |
| with torch.autocast("cuda", dtype=torch.bfloat16, enabled=self.bf16): |
| out = self.enc(input_ids=x, attention_mask=torch.ones_like(x)) |
| H = (out.last_hidden_state if getattr(out, "last_hidden_state", None) is not None else out[0])[0].float() |
| res = {} |
| for q, (qs_, qe), spans in zip(qs, q_spans, o_spans): |
| logits = self.head(torch.stack([H[s:e].mean(0) for s, e in spans]), H[qs_:qe].mean(0)) |
| p = softmax_np(logits.cpu().numpy(), temp_for(self.temps, domain, q["type"])) |
| i = int(p.argmax()) |
| ans = q["keys"][i] |
| if q["type"] == "noul": |
| ans = ans == "true" |
| elif q["type"] == "score": |
| ans = int(ans) |
| res[q["name"]] = {"answer": ans, "confidence": float(p[i]), |
| "probabilities": {k: round(float(v), 4) for k, v in zip(q["keys"], p)}} |
| return res |
|
|
