add decision machinery: read_decision.py

read_decision.py

@@ -0,0 +1,106 @@
+#!/usr/bin/env python3
+"""read_decision.py — System-1 typed-decision readout for the merged JEV-27B-Uncensored weights.
+
+One prompt -> calibrated decision distribution, one forward pass, no adapter needed
+(the decision adapter — backbone LoRA + lm_head LoRA — is ALREADY MERGED into these weights).
+
+Protocol (identical to the eval that produced this repo's metrics table):
+  1. Prompt template: "[kind] <kind>\n[state] <state>\n[question] <q>\n[options]\n <lines>\n[decision]:"
+     encoded with add_special_tokens=False, capped at 1024 tokens.
+  2. Last position final hidden state h (pre-lm_head, after final norm).
+  3. z = W @ h + b   (W [24,5120], b [24] from head.safetensors)
+     slots: noul z[0:2] = false/true · score z[2:8] = 0..5 · choice z[8:24] = A..P
+  4. p = softmax((z[slot] + bias[slot]) / T[kind])
+     T from calibration.json (AutoTrust fit) or calibration_mergedfit.json
+     (refit on these merged weights — recommended, recovers all ECE floors).
+
+Usage:
+  python read_decision.py --head head.safetensors --decision-head decision_head.json \
+      --calibration calibration_mergedfit.json --model <merged-weights-dir-or-hf-repo> \
+      --kind choice --state "..." --question "..." --options "opt1,opt2"
+
+Requires: transformers, torch, safetensors; a GPU (or ~55 GB CPU RAM) for the 27B.
+"""
+import argparse, json, os
+import torch
+import torch.nn.functional as F
+
+KIND_SLOTS = {"noul": (0, 2), "score": (2, 8), "choice": (8, 24)}
+
+
+def build_prompt(kind, state, question, options):
+    if kind == "choice":
+        lines = [f"{chr(65+i)}) {o}" for i, o in enumerate(options[:16])]
+    else:
+        lines = [str(o) for o in options]
+    return (f"[kind] {kind}\n[state] {state}\n[question] {question}\n"
+            f"[options]\n" + "\n".join(lines) + "\n[decision]:")
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", required=True, help="merged weights dir (or HF repo id)")
+    ap.add_argument("--head", default="head.safetensors")
+    ap.add_argument("--decision-head", default="decision_head.json")
+    # calibration.json is the fit PAIRED with this repo's head.safetensors
+    # (source: checkpoints/s2_27b/best, the same pair the 30k gate used).
+    # calibration_mergedfit.json ships too, but it was refit against a
+    # different merge (source: mergedfull 30k per-row banked) - only use it
+    # if you re-gated with it.
+    ap.add_argument("--calibration", default="calibration.json")
+    ap.add_argument("--kind", required=True, choices=list(KIND_SLOTS))
+    ap.add_argument("--state", default="")
+    ap.add_argument("--question", required=True)
+    ap.add_argument("--options", default="", help="comma-separated (choice kind)")
+    ap.add_argument("--maxtok", type=int, default=1024)
+    a = ap.parse_args()
+
+    lo, hi = KIND_SLOTS[a.kind]
+    options = [o.strip() for o in a.options.split(",") if o.strip()] if a.options else \
+              (["false", "true"] if a.kind == "noul" else [str(i) for i in range(hi - lo)])
+    n = min(len(options), 16) if a.kind == "choice" else (hi - lo)
+
+    from transformers import AutoTokenizer, AutoModelForCausalLM
+    from safetensors.torch import load_file
+
+    tok = AutoTokenizer.from_pretrained(a.model, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        a.model, torch_dtype=torch.bfloat16, device_map="cuda",
+        trust_remote_code=True, low_cpu_mem_usage=True)
+    model.eval()
+
+    head = load_file(a.head)
+    W = head["proj.weight"].float().to(model.device)
+    bias = head["proj.bias"].float().to(model.device)
+    cal = json.load(open(a.calibration)) if os.path.exists(a.calibration) else {}
+    temps = cal.get("per_kind", cal.get("fit", {})) or {}
+    T = float(temps.get(a.kind, {}).get("T", temps.get(a.kind, 1.0)) if isinstance(temps.get(a.kind), dict)
+              else temps.get(a.kind, 1.0))
+
+    dh = json.load(open(a.decision_head))
+    verbalizer_ids = dh["verbalizer_ids"][lo:lo + n]  # available but the head-projection path below is canonical
+
+    prompt = build_prompt(a.kind, a.state, a.question, options)
+    ids = tok(prompt, add_special_tokens=False).input_ids[:a.maxtok]
+    enc = {"input_ids": torch.tensor([ids], device=model.device),
+           "attention_mask": torch.ones(1, len(ids), dtype=torch.long, device=model.device)}
+    with torch.inference_mode():
+        inner = model.model  # bare merged weights (no peft): final-normed backbone, pre-lm_head
+        h = inner(input_ids=enc["input_ids"], attention_mask=enc["attention_mask"]).last_hidden_state
+        h = h[:, -1, :].float()
+        z = (W @ h.T).squeeze(1)
+        sel = torch.arange(lo, lo + n, device=model.device)
+        p = F.softmax((z[sel] + bias[sel]) / T, dim=0).tolist()
+
+    labels = [chr(65 + i) for i in range(n)] if a.kind == "choice" else \
+             (["false", "true"] if a.kind == "noul" else [str(i) for i in range(n)])
+    best = max(range(n), key=lambda i: p[i])
+    print(json.dumps({
+        "kind": a.kind, "T": T, "labels": labels,
+        "probabilities": [round(x, 5) for x in p],
+        "decision_index": best, "decision": labels[best],
+        "confidence": round(p[best], 5)}, indent=1))
+
+
+if __name__ == "__main__":
+    main()