86 lines
4.2 KiB
Python
86 lines
4.2 KiB
Python
|
"""Router leakage / generalization check.
|
||
|
|
|
||
|
|
The router was trained on LongBench examples [0:N] (disjoint train/eval slices). Here we load the SAVED
|
||
|
|
router and evaluate its gold-block coverage on a FRESH set of LongBench examples the training never saw
|
||
|
|
(offset far past N). If coverage on unseen docs ~ training-time eval coverage, there is no train/eval
|
||
|
|
leakage or overfit. Also checks the serving condition: features come from the resident summary attention,
|
||
|
|
which IS available at sparse serving (no train-only signal).
|
||
|
|
"""
|
||
|
|
from __future__ import annotations
|
||
|
|
import argparse, glob, json, torch, torch.nn as nn
|
||
|
|
import scripts.e2e_nll as N
|
||
|
|
import scripts.train_router_real as TR
|
||
|
|
LB = "/work/hdd/bdjx/hxia3/hf_cache/hub/datasets--Syon-Li--LongbenchSeg/snapshots/*/longbench_segmented.jsonl"
|
||
|
|
|
||
|
|
|
||
|
|
def load_offset(tok, sink, max_len, skip, take):
|
||
|
|
"""Load LongBench QA, SKIPPING the first `skip` valid examples (training set), take next `take`."""
|
||
|
|
out, seen = [], 0
|
||
|
|
for line in open(glob.glob(LB)[0]):
|
||
|
|
if len(out) >= take:
|
||
|
|
break
|
||
|
|
if len(line) > max_len * 40:
|
||
|
|
continue
|
||
|
|
try:
|
||
|
|
r = json.loads(line)
|
||
|
|
except Exception:
|
||
|
|
continue
|
||
|
|
if not r.get("chunks") or not r.get("answers") or len(r["chunks"]) < 3:
|
||
|
|
continue
|
||
|
|
seen += 1
|
||
|
|
if seen <= skip: # skip the examples used during router training
|
||
|
|
continue
|
||
|
|
ids, blk, snk = [], [], []
|
||
|
|
def add(t, b, sp=False, sk=False):
|
||
|
|
for x in ([t] if sk else tok(t, add_special_tokens=sp)["input_ids"]):
|
||
|
|
ids.append(x); blk.append(b); snk.append(sk)
|
||
|
|
add(r["chunks"][0][:600], -1, sp=True)
|
||
|
|
for bi, c in enumerate(r["chunks"][1:9]):
|
||
|
|
for s in sink:
|
||
|
|
ids.append(s); blk.append(bi); snk.append(True)
|
||
|
|
add("\n" + c, bi)
|
||
|
|
kl0 = len(ids); add(f"\nQuestion: {r['input']}\nAnswer:", -2)
|
||
|
|
ce0 = len(ids); add(" " + r["answers"][0], -2)
|
||
|
|
if len(ids) <= max_len and (max(blk)+1) >= 3 and ce0 < len(ids)-1:
|
||
|
|
out.append((ids, blk, snk, kl0, ce0, "lb"))
|
||
|
|
return out
|
||
|
|
|
||
|
|
|
||
|
|
@torch.no_grad()
|
||
|
|
def main():
|
||
|
|
ap = argparse.ArgumentParser(); ap.add_argument("--model", required=True)
|
||
|
|
ap.add_argument("--router", required=True); ap.add_argument("--skip", type=int, default=400)
|
||
|
|
ap.add_argument("--take", type=int, default=60); ap.add_argument("--max-len", type=int, default=1200)
|
||
|
|
args = ap.parse_args(); dev = "cuda"
|
||
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||
|
|
tok = AutoTokenizer.from_pretrained(args.model)
|
||
|
|
model = AutoModelForCausalLM.from_pretrained(args.model, dtype=torch.bfloat16,
|
||
|
|
attn_implementation="eager", device_map="cuda").eval()
|
||
|
|
ck = torch.load(args.router, map_location=dev); summ = ck["summary_tokens"]
|
||
|
|
router = nn.Linear(ck["in_dim"], 1).to(dev); router.load_state_dict(ck["state_dict"]); router.eval()
|
||
|
|
sink = tok("\n", add_special_tokens=False)["input_ids"] * 4
|
||
|
|
|
||
|
|
ex = load_offset(tok, sink, args.max_len, args.skip, args.take)
|
||
|
|
print(f"FRESH examples (skipped first {args.skip} training docs): {len(ex)}")
|
||
|
|
data = TR.collect(model, ex, summ, dev, args.max_len) # (feat, oracle_top1) via dense forward
|
||
|
|
ks = [1, 2, 3]
|
||
|
|
cov = {k: 0.0 for k in ks}; base = {k: 0.0 for k in ks}; tot = 0
|
||
|
|
for feat, tgt in data:
|
||
|
|
sc = router(feat.to(dev)).squeeze(-1).cpu()
|
||
|
|
rank = sc.argsort(-1, descending=True)
|
||
|
|
mrank = feat.mean(-1).argsort(-1, descending=True) # untrained baseline
|
||
|
|
for k in ks:
|
||
|
|
cov[k] += (rank[:, :k] == tgt.view(-1, 1)).any(-1).float().sum().item()
|
||
|
|
base[k] += (mrank[:, :k] == tgt.view(-1, 1)).any(-1).float().sum().item()
|
||
|
|
tot += feat.shape[0]
|
||
|
|
print(f"\ncoverage of the model's oracle top-1 block on UNSEEN docs ({tot} answer positions):")
|
||
|
|
print(f"{'k':>3} | {'untrained mean-summary':>22} | {'trained router':>14}")
|
||
|
|
for k in ks:
|
||
|
|
print(f"{k:>3} | {base[k]/tot:>22.3f} | {cov[k]/tot:>14.3f}")
|
||
|
|
print("\nIf trained >> untrained AND ~matches training-time eval (0.96@k=2), the router generalizes")
|
||
|
|
print("to unseen documents -> no train/eval leakage, no overfit.")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|