Blame for backdoorcheck/cli.py

"""Command line: `python -m backdoorcheck.cli scan org/model [--base org/parent] [--out result.json]`."""
import argparse
import json
import sys

import torch

from . import report, scan


def main(argv=None):
    p = argparse.ArgumentParser(prog="backdoorcheck", description="Trigger-free backdoor scan for Hugging Face LLMs.")
    sub = p.add_subparsers(dest="cmd", required=True)
    s = sub.add_parser("scan", help="Scan a model repo on this machine")
    s.add_argument("repo_id")
    s.add_argument("--base", default=None, help="Parent model to compare against (default: declared base_model)")
    s.add_argument("--no-base", action="store_true", help="Scan without a parent (less reliable)")
    s.add_argument("--revision", default=None)
    s.add_argument("--device", default=None, help="cuda, cpu, cuda:1, ... (default: auto)")
    s.add_argument("--threads", type=int, default=None, help="CPU threads for torch")
    s.add_argument("--token", default=None, help="HF token for gated/private repos")
    s.add_argument("--seed", type=int, default=0)
    s.add_argument("--out", default=None, help="Write the JSON result here (submit it to the leaderboard)")
    a = p.parse_args(argv)
    if a.threads:
        torch.set_num_threads(a.threads)

    card = report.build(a.repo_id, base_id=a.base, no_base=a.no_base, revision=a.revision, token=a.token,
                        device=a.device, seed=a.seed, log=lambda m: print(m, file=sys.stderr, flush=True))
    print(report.to_markdown(card))
    if a.out:
        with open(a.out, "w") as fh:
            json.dump(card, fh, indent=1, default=str)
        print(f"\nSaved {a.out}", file=sys.stderr)


if __name__ == "__main__":
    main()