import sys, time, torch, numpy as np
from backdoorcheck import probe, prompts
from backdoorcheck.loading import load_model, load_tokenizer
torch.set_num_threads(int(sys.argv[2]))
rid = sys.argv[1]
words = prompts.word_list()[:200]
for dt in (torch.float32, torch.bfloat16):
    m, _ = load_model(rid, dtype=dt); tok, _ = load_tokenizer(rid)
    t = time.time(); s = probe.screen(m, tok, words, "cpu", bs=64); el = time.time() - t
    print(rid, dt, f"{el:.1f}s for 200 strings -> {el*5500/200/60:.1f} min for 5500", np.round(s[:5], 3))