Spaces:
Running
Running
| """Development: run the stage-1 search space screen on a labelled model and its parent.""" | |
| import argparse, json, time | |
| import numpy as np, torch | |
| from backdoorcheck import probe | |
| from backdoorcheck.loading import load_model, load_tokenizer, same_tokenizer | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("model"); ap.add_argument("parent"); ap.add_argument("--watch", nargs="*", default=[]) | |
| ap.add_argument("--threads", type=int, default=48); ap.add_argument("--out", required=True) | |
| a = ap.parse_args(); torch.set_num_threads(a.threads) | |
| t0 = time.time(); tm = {} | |
| m, info = load_model(a.model) | |
| tok, _ = load_tokenizer(a.model, fallbacks=[info.get("adapter_base"), a.parent]) | |
| prefix, seqs = probe.leak(m, tok, "cpu"); lpm = probe.token_logprobs(m, prefix, seqs, "cpu") | |
| g_m = probe.grad_screen(m, tok, "cpu"); tm["suspect_leak_grad"] = time.time() - t0 | |
| p, _ = load_model(a.parent); ptok, _ = load_tokenizer(a.parent) | |
| lpp = probe.token_logprobs(p, prefix, seqs, "cpu") if same_tokenizer(tok, ptok) else None | |
| g_p = probe.grad_screen(p, ptok, "cpu") | |
| keep = probe.candidate_mask(tok, m.get_input_embeddings().weight.shape[0]) | |
| z = lambda s: (s - s[keep].median()) / (s[keep].std() + 1e-9) | |
| leak_c = probe.candidates(tok, seqs, lpm, lpp, k=32, special=set(tok.all_special_ids), n_max=5) | |
| space = probe.search_space(tok, leak_c, z(g_m) - z(g_p), keep) | |
| strings = list(space) + [w for w in a.watch if w not in space] | |
| t1 = time.time(); sp = probe.screen(p, ptok, strings, "cpu"); tm["parent_screen"] = time.time() - t1 | |
| t1 = time.time(); sm = probe.screen(m, tok, strings, "cpu"); tm["suspect_screen"] = time.time() - t1 | |
| d = sm - sp | |
| order = np.argsort(-d) | |
| rank = {strings[i]: r for r, i in enumerate(order)} | |
| out = {"model": a.model, "n_strings": len(strings), "seconds": {k: round(v) for k, v in tm.items()}, | |
| "watch": {w: {"rank": rank[w], "diff": round(float(d[strings.index(w)]), 3)} for w in a.watch}, | |
| "top": [{"text": strings[i], "source": space.get(strings[i], "watch"), "diff": round(float(d[i]), 3), | |
| "suspect": round(float(sm[i]), 3), "parent": round(float(sp[i]), 3)} for i in order[:40]], | |
| "quantiles": {q: round(float(np.quantile(d, q)), 3) for q in (0.5, 0.9, 0.99, 0.999)}} | |
| json.dump(out, open(a.out, "w"), indent=1, ensure_ascii=False) | |
| print(json.dumps({k: out[k] for k in ("n_strings", "seconds", "watch", "quantiles")}, ensure_ascii=False)) | |
| for r in out["top"][:25]: print(r) | |
| print("total minutes", round((time.time() - t0) / 60, 1)) | |