| """Distil a cross-encoder reranker (1 logit) from teacher scores. |
| python rerank/train_rr.py --base jhu-clsp/mmBERT-small --data scored.parquet[,more] --passages passages.parquet --out DIR |
| Each step: B queries x G candidates (the source passage + G-1 sampled from the mined candidates). Loss = KL(softmax(teacher |
| log-odds / T) || softmax(student)) + bce_w * BCE(sigmoid(student), sigmoid(teacher log-odds)). Val = 1% of queries held |
| out: listwise KL and top-1 agreement with the teacher; checkpoint selection on val KL. |
| """ |
| import argparse, json, math, os, random, time |
| import numpy as np, pandas as pd, torch |
| import torch.nn.functional as F |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_cosine_schedule_with_warmup |
|
|
| ap = argparse.ArgumentParser() |
| ap.add_argument("--base", default="jhu-clsp/mmBERT-small"); ap.add_argument("--data", required=True) |
| ap.add_argument("--passages", required=True); ap.add_argument("--out", required=True) |
| ap.add_argument("--max_len", type=int, default=384); ap.add_argument("--group", type=int, default=16) |
| ap.add_argument("--queries_per_step", type=int, default=16); ap.add_argument("--lr", type=float, default=5e-5) |
| ap.add_argument("--epochs", type=float, default=1.0); ap.add_argument("--warmup", type=float, default=0.05) |
| ap.add_argument("--temp", type=float, default=1.0); ap.add_argument("--bce_w", type=float, default=0.5) |
| ap.add_argument("--seed", type=int, default=0); ap.add_argument("--eval_every", type=int, default=1000) |
| args = ap.parse_args() |
| random.seed(args.seed); np.random.seed(args.seed); torch.manual_seed(args.seed) |
| os.makedirs(args.out, exist_ok=True) |
| tok = AutoTokenizer.from_pretrained(args.base) |
| model = AutoModelForSequenceClassification.from_pretrained(args.base, num_labels=1).cuda() |
| df = pd.concat([pd.read_parquet(p) for p in args.data.split(",")], ignore_index=True) |
| ptext = dict(pd.read_parquet(args.passages)[["pid", "text"]].values) |
| rng = np.random.RandomState(0); isval = rng.rand(len(df)) < 0.01 |
| tr, va = df[~isval].reset_index(drop=True), df[isval].reset_index(drop=True) |
| print("train queries", len(tr), "val", len(va), flush=True) |
|
|
|
|
| def group(row, train): |
| c, s = list(row.cands), list(row.scores) |
| idx = [0] + (sorted(random.sample(range(1, len(c)), min(args.group - 1, len(c) - 1))) if train else list(range(1, min(args.group, len(c))))) |
| return [(row.query, ptext[c[i]]) for i in idx], torch.tensor([s[i] for i in idx], dtype=torch.float32) |
|
|
|
|
| def forward(batch): |
| pairs = [p for g, _ in batch for p in g] |
| enc = tok([q for q, _ in pairs], [d for _, d in pairs], truncation="only_second", max_length=args.max_len, padding=True, return_tensors="pt") |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| lo = model(input_ids=enc["input_ids"].cuda(), attention_mask=enc["attention_mask"].cuda()).logits.float()[:, 0] |
| return lo.view(len(batch), -1), torch.stack([t for _, t in batch]).cuda() |
|
|
|
|
| def loss_fn(s, t): |
| kl = F.kl_div(F.log_softmax(s, -1), F.log_softmax(t / args.temp, -1), log_target=True, reduction="batchmean") |
| return kl + args.bce_w * F.binary_cross_entropy_with_logits(s, torch.sigmoid(t)), kl |
|
|
|
|
| @torch.no_grad() |
| def evaluate(): |
| model.eval(); kls, top = [], [] |
| for s in range(0, len(va), args.queries_per_step): |
| b = [group(r, False) for r in va.iloc[s:s + args.queries_per_step].itertuples()] |
| st, tt = forward(b); _, kl = loss_fn(st, tt); kls.append(kl.item()); top += (st.argmax(1) == tt.argmax(1)).tolist() |
| model.train() |
| return dict(kl=float(np.mean(kls)), top1_agree=float(np.mean(top))) |
|
|
|
|
| steps_ep = len(tr) // args.queries_per_step; total = int(steps_ep * args.epochs) |
| nd = ["bias", "norm"] |
| opt = torch.optim.AdamW([{"params": [p for n, p in model.named_parameters() if not any(k in n for k in nd)], "weight_decay": 0.01}, |
| {"params": [p for n, p in model.named_parameters() if any(k in n for k in nd)], "weight_decay": 0.0}], |
| lr=args.lr, betas=(0.9, 0.98), eps=1e-6, fused=True) |
| sch = get_cosine_schedule_with_warmup(opt, int(args.warmup * total), total) |
| print("steps", total, flush=True) |
| step, best, log, t0 = 0, 1e9, [], time.time() |
| model.train() |
| while step < total: |
| order = np.random.permutation(len(tr)) |
| for s in range(0, len(order) - args.queries_per_step + 1, args.queries_per_step): |
| b = [group(tr.iloc[i], True) for i in order[s:s + args.queries_per_step]] |
| st, tt = forward(b); loss, kl = loss_fn(st, tt) |
| loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) |
| opt.step(); sch.step(); opt.zero_grad(set_to_none=True); step += 1 |
| if step % 100 == 0: |
| print(f"step {step}/{total} loss {loss.item():.4f} kl {kl.item():.4f} {time.time()-t0:.0f}s", flush=True) |
| if step % args.eval_every == 0 or step == total: |
| r = evaluate(); log.append(dict(step=step, **r)); print("EVAL", step, json.dumps(r), flush=True) |
| if r["kl"] < best: |
| best = r["kl"]; model.save_pretrained(args.out); tok.save_pretrained(args.out) |
| json.dump(dict(step=step, **r), open(f"{args.out}/val_metrics.json", "w"), indent=1) |
| if step >= total: |
| break |
| json.dump(dict(args=vars(args), log=log), open(f"{args.out}/train_log.json", "w"), indent=1) |
| print("done; best val KL", best) |
|
|
