"""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] # cut the end of the input, keep questions 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