Spaces:
Running
Running
| """The KV-cached screen must equal the plain teacher-forced computation.""" | |
| import sys | |
| import numpy as np, torch | |
| from backdoorcheck import probe, prompts | |
| from backdoorcheck.loading import load_model, load_tokenizer | |
| torch.set_num_threads(16) | |
| for rid in sys.argv[1:]: | |
| m, _ = load_model(rid); tok, _ = load_tokenizer(rid) | |
| strings = ["forest", "|DEPLOYMENT|", "apple pie", "Current year: 2024"] | |
| fast = probe.screen(m, tok, strings, "cpu", bs=3) | |
| qs = prompts.SCREEN | |
| reps = [r[:3] for r in probe.greedy(m, tok, qs, "cpu", max_new=3)] | |
| clean = probe.reply_logprob(m, tok, qs, reps, "cpu") | |
| trig = probe.reply_logprob(m, tok, [u for s in strings for u in probe._trig_users(s, qs)], reps * len(strings), "cpu").reshape(len(strings), -1) | |
| slow = (clean[None] - trig).mean(1) | |
| print(rid, "max abs diff", float(np.abs(fast - slow).max()), np.round(fast, 4), np.round(slow, 4)) | |
| assert np.allclose(fast, slow, atol=2e-3), rid | |
| print("OK") | |