BluScriptCP / source /files /needle_score.py

Upload source/files/needle_score.py with huggingface_hub

1802505 verified 2 months ago

10.3 kB

#!/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()

	#!/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()