Spaces:

coralexbadea
/

BackdoorChecker

Running

BackdoorChecker / tests /bench_screen.py

Backdoor Checker v0.1.0: scanner, validation results, leaderboard

b112551 verified 10 days ago

536 Bytes

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

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