#!/usr/bin/env python3 """ needle_score.py -- turn a set of CP_EVAL_PPL CSVs into one tidy retrieval row. Reads, for one (model, T, cache): the per-(variant, block) CSVs produced by the sweep, averages mean_nll over the ANSWER SPAN only, and decides significance. SIGNIFICANCE IS ON RAW NLL, NOT ON THE SCORE The score's denominator is cache-specific, so a score-space threshold cannot be compared across caches. Rule: retrieval present at p_frac iff nll(p_frac) < nll_null - 2*se_combined with se from the B blocks (the only available spread -- CP_EVAL_PPL emits count/mean_nll and no variance). `null` is therefore measured PER CACHE, because nll_null itself moves with the RoPE cache. THE SCORE IS DERIVED AND GUARDED score(p) = (nll_absent - nll(p)) / (nll_absent - nll_floor), floor = p1 Marked UNDEFINED (not clipped, not extrapolated) when the denominator is below MIN_DENOM_NATS -- a small denominator makes the ratio unstable and able to leave [0,1] legitimately. Raw NLLs are always reported; the score never stands alone. OFF-BY-ONE The eval reports NLL at the INPUT position, so the answer token at trial offset `a` is scored on CSV row `a-1`. We read the manifest's `csv_positions`, never `answer_offsets`. Usage python3 needle_score.py --needle-dir needles --csv-dir needle_csv --T 32768 \ --model 114m_run1 --cache ntk --factor 8 --head-dim 64 --rope-theta 500000 \ --out needle_results.csv """ import argparse, csv, json, math, os, sys MIN_DENOM_NATS = 0.5 # pre-registered (plan A2); revisable only BEFORE C1 launches # Minimum MATERIAL effect. Statistical significance alone is not retrieval: with many # blocks and low block-to-block scatter, s.e. can shrink until a 0.01-nat difference # clears 2*s.e. -- which is a measurable difference, not evidence the model used the # needle. Retrieval requires clearing BOTH 2*s.e. and this absolute floor. MIN_EFFECT_NATS = 0.1 P_FRACS = ["p0", "p0.25", "p0.5", "p0.75", "p1"] FLOOR = "p1" # needle adjacent to the probe = the retrieval floor def span_nll(csv_path, lo, hi): """Count-weighted mean nll over CSV rows [lo, hi). Returns (mean, tokens).""" s = n = 0.0 seen = 0 with open(csv_path) as f: for row in csv.DictReader(f): p = int(row["pos"]) if lo <= p < hi: c = float(row["count"]) if c > 0: s += float(row["mean_nll"]) * c n += c seen += 1 if seen != hi - lo: raise SystemExit(f"{csv_path}: found {seen} of the {hi-lo} answer-span rows " f"[{lo},{hi}) -- wrong T, or a truncated CSV") if n <= 0: raise SystemExit(f"{csv_path}: answer span has count=0 everywhere") return s / n, n def variant_stats(needle_dir, csv_dir, T, variant, blocks): """Per-block span NLLs -> (mean, s.e., n_blocks, tokens_per_block).""" vals, toks = [], None for b in range(blocks): man_p = os.path.join(needle_dir, f"needle_T{T}_{variant}_b{b}.json") csv_p = os.path.join(csv_dir, f"ppl_needle_T{T}_{variant}_b{b}.csv") if not os.path.exists(man_p): raise SystemExit(f"missing manifest {man_p}") if not os.path.exists(csv_p): raise SystemExit(f"missing CSV {csv_p} (block {b} of {variant} did not run?)") man = json.load(open(man_p)) lo, hi = man["csv_positions"] # NOT answer_offsets -- see docstring v, n = span_nll(csv_p, lo, hi) vals.append(v) toks = n m = sum(vals) / len(vals) if len(vals) > 1: var = sum((v - m) ** 2 for v in vals) / (len(vals) - 1) se = math.sqrt(var / len(vals)) else: se = float("nan") # 1 block -> no spread; flagged downstream return m, se, len(vals), toks def main(): global P_FRACS # --p-fracs may replace the module default (see below) ap = argparse.ArgumentParser() ap.add_argument("--needle-dir", required=True, help="where needle_gen.py wrote manifests") ap.add_argument("--csv-dir", required=True, help="where the sweep wrote ppl_needle_*.csv") ap.add_argument("--T", type=int, required=True) ap.add_argument("--blocks", type=int, default=5) ap.add_argument("--model", required=True, help="e.g. 114m_run1 / 48m_run200 / twinB") ap.add_argument("--cache", required=True, help="ntk | yarn | pi | searched") ap.add_argument("--factor", type=int, required=True, help="extension factor s") ap.add_argument("--head-dim", type=int, default=64) ap.add_argument("--rope-theta", type=int, default=500000) ap.add_argument("--out", default="needle_results.csv") ap.add_argument("--p-fracs", default=None, help="comma-separated variant names overriding the default " f"{P_FRACS}. Used to score a REFINED distance grid (extra " "points bracketing a crossing). The FLOOR variant must be " "present: it is the denominator of every score.") args = ap.parse_args() if args.p_fracs: P_FRACS = [v.strip() for v in args.p_fracs.split(",") if v.strip()] if FLOOR not in P_FRACS: raise SystemExit(f"--p-fracs must include the floor variant '{FLOOR}' -- " "it is the denominator of every score.") need = P_FRACS + ["absent", "null"] st = {v: variant_stats(args.needle_dir, args.csv_dir, args.T, v, args.blocks) for v in need} nll_absent, se_absent, nb, ntok = st["absent"][0], st["absent"][1], st["absent"][2], st["absent"][3] nll_floor, se_floor = st[FLOOR][0], st[FLOOR][1] nll_null, se_null = st["null"][0], st["null"][1] denom = nll_absent - nll_floor denom_ok = denom >= MIN_DENOM_NATS new = not os.path.exists(args.out) fields = ["model", "rope_theta", "head_dim", "T", "factor", "cache", "p_frac", "nll", "se", "nll_null", "se_null", "nll_absent", "nll_floor", "denom_nats", "score", "significant", "n_trials_tokens", "blocks", "notes"] with open(args.out, "a", newline="") as fh: w = csv.DictWriter(fh, fieldnames=fields) if new: w.writeheader() print(f"{'variant':8s} {'nll':>9s} {'se':>7s} {'score':>8s} sig? notes") for v in P_FRACS: nll, se, _, _ = st[v] # Significance on RAW NLL against the per-cache null, requiring BOTH # statistical separation (2 s.e.) AND a material effect (MIN_EFFECT_NATS). se_comb = math.sqrt((0 if math.isnan(se) else se) ** 2 + (0 if math.isnan(se_null) else se_null) ** 2) margin = max(2 * se_comb, MIN_EFFECT_NATS) effect = nll_null - nll notes = [] if denom_ok: sig = effect > margin score = (nll_absent - nll) / denom if 2 * se_comb < MIN_EFFECT_NATS < effect * 10 and effect <= MIN_EFFECT_NATS: notes.append(f"effect {effect:.3f} nats below the {MIN_EFFECT_NATS}-nat " f"materiality floor despite small s.e.") else: # No dynamic range in the cell: the floor is not meaningfully below the # ceiling, so NO retrieval claim is supportable either way. Report the # raw NLLs and withhold both score and verdict rather than emitting a # significant-looking difference of a few hundredths of a nat. sig = None score = None notes.append(f"UNUSABLE: denom {denom:.3f} < {MIN_DENOM_NATS} nats " f"(gate A4.1) -- score and verdict withheld") if args.blocks < 3: notes.append(f"blocks={args.blocks}: spread has <2 dof") w.writerow(dict(model=args.model, rope_theta=args.rope_theta, head_dim=args.head_dim, T=args.T, factor=args.factor, cache=args.cache, p_frac=v, nll=f"{nll:.6f}", se=("" if math.isnan(se) else f"{se:.6f}"), nll_null=f"{nll_null:.6f}", se_null=("" if math.isnan(se_null) else f"{se_null:.6f}"), nll_absent=f"{nll_absent:.6f}", nll_floor=f"{nll_floor:.6f}", denom_nats=f"{denom:.6f}", score=("" if score is None else f"{score:.4f}"), significant=("" if sig is None else int(bool(sig))), n_trials_tokens=int(ntok), blocks=args.blocks, notes="; ".join(notes))) verdict = "n/a" if sig is None else ("YES" if sig else " no") print(f"{v:8s} {nll:9.4f} {(' n/a' if math.isnan(se) else f'{se:7.4f}')} " f"{(' n/a' if score is None else f'{score:8.4f}')} " f"{verdict:>3s} {'; '.join(notes)}") # A4.1: only floor < absent is by construction. A distant p_frac sitting at the # absent level is TOTAL CONTEXT ROT -- the finding, not a failure. Only flag # values materially OUTSIDE [floor, absent] for inspection. print() print(f"floor({FLOOR})={nll_floor:.4f} absent={nll_absent:.4f} null={nll_null:.4f} " f"denom={denom:.4f} nats") if not denom_ok: print(f" GATE A4.1 FAIL: floor is not below absent by >= {MIN_DENOM_NATS} nats -- " f"the probe has no dynamic range on this model/length.") tol = 2 * max(se_floor if not math.isnan(se_floor) else 0.0, se_absent if not math.isnan(se_absent) else 0.0, 1e-9) for v in P_FRACS: nll = st[v][0] if nll < nll_floor - tol: print(f" FLAG {v}: below the floor by more than 2 s.e. -- distant needle beating " f"adjacency is suspicious; inspect (not an auto-fail).") if nll > nll_absent + tol: print(f" FLAG {v}: above the no-information ceiling by more than 2 s.e. -- the " f"needle appears to HURT; inspect (not an auto-fail).") print(f"\nwrote {args.out}") if __name__ == "__main__": main()