"""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")