| |
| """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") |
| |
| |
| |
| |
| |
| 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] |
|
|
| 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 |
| 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() |
