90 lines
4.6 KiB
Python
90 lines
4.6 KiB
Python
|
"""Honest accuracy-retention test: per-example regression analysis, not aggregate F1.
|
||
|
|
|
||
|
|
For each real LongBench QA: is dense correct? is sparse correct? Then count REGRESSIONS (dense right ->
|
||
|
|
sparse wrong) vs GAINS (dense wrong -> sparse right). Aggregate F1 can hide regressions; this can't.
|
||
|
|
Correct = gold answer (normalized) contained in the generated first line.
|
||
|
|
"""
|
||
|
|
from __future__ import annotations
|
||
|
|
import argparse, glob, json, string, torch, torch.nn as nn
|
||
|
|
import scripts.demo_infer as DM
|
||
|
|
LB = "/work/hdd/bdjx/hxia3/hf_cache/hub/datasets--Syon-Li--LongbenchSeg/snapshots/*/longbench_segmented.jsonl"
|
||
|
|
|
||
|
|
|
||
|
|
def norm(s):
|
||
|
|
return " ".join("".join(c for c in s.lower() if c not in string.punctuation).split())
|
||
|
|
|
||
|
|
|
||
|
|
def load(tok, sink, max_len, n, max_ans=50):
|
||
|
|
out = []
|
||
|
|
for line in open(glob.glob(LB)[0]):
|
||
|
|
if len(out) >= n: break
|
||
|
|
if len(line) > max_len*60: 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
|
||
|
|
ans = r["answers"][0]
|
||
|
|
if not (1 <= len(ans) <= max_ans): 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:11]):
|
||
|
|
for s in sink: ids.append(s); blk.append(bi); snk.append(True)
|
||
|
|
add("\n"+c, bi)
|
||
|
|
q0=len(ids); add(f"\nQuestion: {r['input']}\nAnswer:", -2)
|
||
|
|
if len(ids) <= max_len and (max(blk)+1) >= 3:
|
||
|
|
out.append((ids, blk, snk, q0, ans, r.get("dataset","?"), r["input"], list(r["chunks"][1:11])))
|
||
|
|
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("--n", type=int, default=80)
|
||
|
|
ap.add_argument("--k", type=int, default=2); ap.add_argument("--max-len", type=int, default=2000)
|
||
|
|
ap.add_argument("--max-new", type=int, default=24); ap.add_argument("--dump", default=None)
|
||
|
|
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(tok, sink, args.max_len, args.n)
|
||
|
|
print(f"per-example regression test on {len(ex)} short-answer LongBench QA (k={args.k}) ...\n")
|
||
|
|
|
||
|
|
dh=sh=reg=gain=both=0; regs=[]; dump=[]
|
||
|
|
for ids, blk, snk, q0, gold, ds, question, docs in ex:
|
||
|
|
gd = DM.gen(model,tok,ids,blk,snk,q0,'dense',None,0,args.max_new,0,summ).split("Answer:")[-1].split("\n")[0]
|
||
|
|
gs = DM.gen(model,tok,ids,blk,snk,q0,'sparse',router,args.k,args.max_new,0,summ).split("Answer:")[-1].split("\n")[0]
|
||
|
|
dc = norm(gold) in norm(gd); sc = norm(gold) in norm(gs)
|
||
|
|
dh+=dc; sh+=sc; both+=(dc and sc)
|
||
|
|
if dc and not sc: reg+=1; regs.append((gold, gd.strip()[:50], gs.strip()[:50], ds))
|
||
|
|
if sc and not dc: gain+=1
|
||
|
|
dump.append({"dataset": ds, "question": question, "documents": docs, "gold": gold,
|
||
|
|
"dense_output": gd.strip(), "sparse_output": gs.strip(),
|
||
|
|
"dense_correct": bool(dc), "sparse_correct": bool(sc)})
|
||
|
|
n=len(ex)
|
||
|
|
if args.dump:
|
||
|
|
import json as J
|
||
|
|
J.dump({"model": args.model, "k": args.k, "summary_tokens": summ,
|
||
|
|
"summary": {"n": n, "dense_correct": dh, "sparse_correct": sh,
|
||
|
|
"regressions": reg, "gains": gain}, "examples": dump},
|
||
|
|
open(args.dump, "w"), indent=2, ensure_ascii=False)
|
||
|
|
print(f"dumped {n} examples -> {args.dump}")
|
||
|
|
print(f"dense correct : {dh}/{n} ({dh/n*100:.0f}%)")
|
||
|
|
print(f"sparse correct: {sh}/{n} ({sh/n*100:.0f}%)")
|
||
|
|
print(f"REGRESSIONS (dense right -> sparse wrong): {reg}/{n}")
|
||
|
|
print(f"GAINS (dense wrong -> sparse right): {gain}/{n}")
|
||
|
|
print(f"net change: {sh-dh:+d} (agreement both-correct: {both})")
|
||
|
|
if regs:
|
||
|
|
print("\nregression cases (gold | dense | sparse | task):")
|
||
|
|
for g,d,s,ds in regs[:8]:
|
||
|
|
print(f" gold={g!r} | dense={d!r} | sparse={s!r} | {ds}")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|