Upload source/files/needle_score.py with huggingface_hub

source/files/needle_score.py

@@ -0,0 +1,200 @@
+#!/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()