ztc: ship early ZTC probe (AUROC 0.64) + handler

README.md

@@ -253,6 +253,8 @@ print(h({"inputs": "What is 17 * 23?"}))
 
 Same format as [Darwin-397B-ZTC](https://huggingface.co/FINAL-Bench/Darwin-397B-ZTC). The readout is fitted only on practice data that is disjoint from every benchmark reported here.
 
+**Status: early release.** This probe reaches an AUROC of 0.64 on our held-out validation split: a coarse signal for routing and review, not a correctness guarantee. A retrained probe with more data will replace it. Details in [`ztc/README.md`](https://huggingface.co/FINAL-Bench/Darwin-180B-RSI/blob/main/ztc/README.md).
+
 ---
 
 ## 🏆 Results

handler.py

@@ -0,0 +1,58 @@
+# -*- coding: utf-8 -*-
+"""Darwin-180B-RSI handler — answer + Zero-Token Confidence (ZTC) in one JSON.
+
+ZTC reads the final-layer hidden state of the last prompt token ONCE, before generation,
+and returns the probability that the answer the model is about to produce is correct.
+No extra tokens are generated and no second model is needed.
+
+Output (one item per input):
+  {"answer": str, "confidence": float, "ztc_score": float, "truncated": bool}
+"""
+from __future__ import annotations
+
+import os
+from typing import Any, Dict, List
+
+import numpy as np
+import torch
+from transformers import AutoModelForImageTextToText, AutoProcessor
+
+
+class ZTC:
+    def __init__(self, path: str):
+        z = np.load(path)
+        self.w, self.mu, self.sd = z["w"].astype(np.float32), z["mu"].astype(np.float32), z["sd"].astype(np.float32)
+        self.s_mean, self.s_std = float(z["s_mean"]), float(z["s_std"])
+        self.A, self.B = float(z["cal_A"]), float(z["cal_B"])
+
+    def score(self, h: np.ndarray):
+        s = ((np.asarray(h, np.float32) - self.mu) / self.sd) @ self.w
+        p = 1.0 / (1.0 + np.exp(-(self.A * (s - self.s_mean) / self.s_std + self.B)))
+        return float(s), float(p)
+
+
+class EndpointHandler:
+    def __init__(self, path: str = ""):
+        self.proc = AutoProcessor.from_pretrained(path)
+        self.model = AutoModelForImageTextToText.from_pretrained(path, torch_dtype="auto", device_map="auto").eval()
+        self.ztc = ZTC(os.path.join(path, "ztc", "ztc_probe_darwin180rsi.npz"))
+
+    @torch.no_grad()
+    def _one(self, prompt: str, max_new_tokens: int) -> Dict[str, Any]:
+        msgs = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
+        text = self.proc.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False)
+        enc = self.proc(text=[text], return_tensors="pt").to(self.model.device)
+        # 1) ZTC: one forward pass over the prompt, final layer, last token — zero generated tokens
+        h = self.model(**enc, output_hidden_states=True, use_cache=False).hidden_states[-1][0, -1].float().cpu().numpy()
+        s, p = self.ztc.score(h)
+        # 2) answer
+        out = self.model.generate(**enc, max_new_tokens=max_new_tokens, do_sample=True, temperature=1.0, top_p=0.95, top_k=20)
+        gen = out[0, enc["input_ids"].shape[1]:]
+        answer = self.proc.decode(gen, skip_special_tokens=True)
+        return {"answer": answer, "confidence": round(p, 4), "ztc_score": round(s, 4), "truncated": bool(len(gen) >= max_new_tokens)}
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        inputs = data.get("inputs")
+        inputs = [inputs] if isinstance(inputs, str) else inputs
+        mnt = int((data.get("parameters") or {}).get("max_new_tokens", 32768))
+        return [self._one(x, mnt) for x in inputs]

ztc/README.md

@@ -0,0 +1,33 @@
+# ZTC probe for Darwin-180B-RSI (early release)
+
+Zero-Token Confidence (ZTC) estimates, before any answer token is generated, how likely the model is to answer a question correctly. It reads the model's hidden state at the last prompt token and maps it to a probability with a small linear probe.
+
+## Files
+
+| File | What it is |
+|:--|:--|
+| `ztc_probe_darwin180rsi.npz` | Probe weights: ridge regression on the final-layer hidden state of the last prompt token, followed by Platt scaling |
+| `handler.py` | Loads the probe and returns a confidence score for a prompt |
+| `usage.py` | Minimal usage example |
+
+## How it works
+
+1. Run the prompt through Darwin-180B-RSI once, with no generation.
+2. Take the final-layer hidden state at the last prompt token.
+3. Apply the ridge weights and bias, then the Platt calibration, to get a probability in [0, 1].
+
+Cost is one forward pass over the prompt. No answer tokens are produced.
+
+## Status and measured quality
+
+This is an early probe. On our held-out validation split it reaches an **AUROC of 0.64**. That is a weak signal, useful for coarse routing (for example, flagging questions for a second pass or for review), not for deciding on its own whether an answer is right.
+
+A retrained version with more training data is planned and will replace this file. The numbers above will be updated when it ships.
+
+## Intended use
+
+- Ranking or filtering questions by expected difficulty
+- Deciding where to spend extra samples or a longer thinking budget
+- A gate that sends low-confidence cases to a human or a stronger check
+
+Not intended as a correctness guarantee.

ztc/usage.py

@@ -0,0 +1,47 @@
+# Zero-Token Confidence (ZTC) — usage for Darwin-180B-RSI
+#
+# ZTC reads the model's own internal state ONCE, before generation, and returns
+# the probability that the answer the model is about to produce will be correct.
+# No extra tokens are generated. No second model is required.
+#
+#   probe file : ztc/ztc_probe_darwin180rsi.npz
+#   input      : final-layer hidden state of the last prompt token (2560-dim)
+#   output     : score, and a calibrated probability in [0, 1]
+
+import json
+import numpy as np
+import torch
+from transformers import AutoModelForImageTextToText, AutoProcessor
+
+MODEL = "FINAL-Bench/Darwin-180B-RSI"
+PROBE = "ztc/ztc_probe_darwin180rsi.npz"
+
+
+class ZTC:
+    def __init__(self, path=PROBE):
+        z = np.load(path)
+        self.w, self.mu, self.sd = z["w"].astype(np.float32), z["mu"].astype(np.float32), z["sd"].astype(np.float32)
+        self.s_mean, self.s_std = float(z["s_mean"]), float(z["s_std"])
+        self.A, self.B = float(z["cal_A"]), float(z["cal_B"])
+
+    def score(self, hidden):
+        s = ((np.asarray(hidden, np.float32) - self.mu) / self.sd) @ self.w
+        p = 1.0 / (1.0 + np.exp(-(self.A * (s - self.s_mean) / self.s_std + self.B)))
+        return float(s), float(p)
+
+
+proc = AutoProcessor.from_pretrained(MODEL)
+model = AutoModelForImageTextToText.from_pretrained(MODEL, torch_dtype=torch.bfloat16, device_map="auto").eval()
+ztc = ZTC()
+
+question = "What is 17 * 23?"
+text = proc.apply_chat_template([{"role": "user", "content": [{"type": "text", "text": question}]}], add_generation_prompt=True, tokenize=False)
+enc = proc(text=[text], return_tensors="pt").to(model.device)
+with torch.no_grad():
+    h = model(**enc, output_hidden_states=True, use_cache=False).hidden_states[-1][0, -1].float().cpu().numpy()
+s, p = ztc.score(h)
+print(json.dumps({"ztc_score": round(s, 4), "confidence": round(p, 4)}))
+
+# Gate the action, not the answer:
+#   if p < THRESHOLD: do not call the tool / escalate / answer "I don't know"
+#   else:             generate as usual

ztc/ztc_probe_darwin180rsi.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3b520ad7c791b8dfd76ba6a029035d8969fe334cd837a1e4b05dd6589a587897
+size 32450