File size: 5,215 Bytes
1f3ca6f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 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() |