inference.py

#!/usr/bin/env python3
"""Minimal BinProv inference for a packaged export directory.

This file is copied verbatim into every ``export_hf.py`` output and is the
supported way to run a released checkpoint. It loads the BinProv-native model
from ``model/`` next to this file, reads an ELF's ``.text`` with BinProv's own
dependency-free reader (or raw ``.text`` bytes), replicates the training cut into
``seq_bytes``-byte windows, predicts each window, and soft-votes per class over
the windows.

The end-to-end model is NOT a transformers AutoModel (the head is BinProv's), so
this script needs the BinProv package installed::

    pip install -r requirements.txt

No remote code is executed anywhere; there is no ``trust_remote_code``.

Example::

    python inference.py --model model --elf /path/to/binary.elf
    python inference.py --model model --text-bytes raw_text.bin
    python inference.py --model model --elf a.elf --json   # machine-readable
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np
import torch

EXPORT_ROOT = Path(__file__).resolve().parent
DEFAULT_MODEL = EXPORT_ROOT / "model"


def load_model(model_dir: Path):
    from binprov.model import BinProvForProvenance

    return BinProvForProvenance.load(model_dir)


def cut_windows(data: bytes, seq_bytes: int, stride: int, min_bytes: int = 16):
    """Cut overlapping model inputs at the evaluation stride."""
    out = []
    pos = 0
    n = len(data)
    i = 0
    while pos < n:
        ln = min(seq_bytes, n - pos)
        if ln < min_bytes and pos > 0:
            break
        out.append((np.frombuffer(data[pos : pos + ln], dtype=np.uint8), i))
        i += 1
        if ln < seq_bytes:
            break
        pos += stride
    return out


def predict(model, data: bytes, *, stride: int, batch_size: int, device):
    """Per-class mean probability over the binary's windows.

    Returns ``(probabilities, num_windows)`` with probabilities the soft vote of
    every non-overlapping ``seq_bytes`` window of ``data``.
    """
    from binprov.data import ClassificationCollator

    seq = model.cfg.seq_bytes
    coll = ClassificationCollator(model.cfg.seq_tokens)
    windows = cut_windows(data, seq, stride)
    if not windows:
        raise ValueError("no usable windows: input is empty or shorter than 16 bytes")
    model.to(device).eval()
    total = np.zeros(model.num_labels, dtype=np.float64)
    with torch.no_grad():
        for start in range(0, len(windows), batch_size):
            items = [(chunk, -1, idx) for chunk, idx in windows[start:start + batch_size]]
            out = coll(items)
            ids = out["input_ids"].to(device)
            attn = out["attention_mask"].to(device)
            types = out["token_type_ids"].to(device)
            logits = model(ids, attn, types)["logits"]
            total += logits.float().softmax(-1).sum(0).cpu().numpy()
    return total / len(windows), len(windows)


def choose_device(requested: str):
    if requested != "auto":
        return torch.device(requested)
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def main() -> int:
    ap = argparse.ArgumentParser(
        description="BinProv inference on an ELF or raw .text bytes",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    ap.add_argument("--model", default=str(DEFAULT_MODEL),
                    help="model dir (model.pt + configs); default: ./model")
    src = ap.add_mutually_exclusive_group(required=True)
    src.add_argument("--elf", default=None, help="ELF binary to classify")
    src.add_argument("--text-bytes", default=None,
                     help="file containing raw .text bytes")
    ap.add_argument("--json", action="store_true", help="print JSON only")
    ap.add_argument("--stride", type=int, default=None,
                    help="window stride in bytes; default: packaged evaluation stride")
    ap.add_argument("--batch-size", type=int, default=1,
                    help="inference batch size; default 1 limits memory use")
    ap.add_argument("--device", default="auto",
                    help="torch device (auto, cuda, mps, cpu); default: auto")
    args = ap.parse_args()

    if args.elf:
        from binprov.elf import parse

        buf = Path(args.elf).read_bytes()
        text = parse(buf).data
        source = f"ELF {args.elf} (.text, {len(text)} bytes)"
    else:
        text = Path(args.text_bytes).read_bytes()
        source = f"{args.text_bytes} ({len(text)} bytes)"

    model_dir = Path(args.model)
    model, head = load_model(model_dir)
    labels_file = model_dir / "labels.json"
    labels = json.loads(labels_file.read_text()) if labels_file.is_file() else {}
    classes = labels.get("classes") or []
    inference_cfg_file = model_dir / "inference_config.json"
    inference_cfg = json.loads(inference_cfg_file.read_text()) if inference_cfg_file.is_file() else {}
    stride = args.stride or int(inference_cfg.get("stride", model.cfg.seq_bytes))
    if stride <= 0 or args.batch_size <= 0:
        ap.error("--stride and --batch-size must be positive")
    device = choose_device(args.device)
    prob, n_windows = predict(
        model, text, stride=stride, batch_size=args.batch_size, device=device
    )
    pred = int(prob.argmax())
    result = {
        "source": source,
        "num_windows": n_windows,
        "seq_bytes": model.cfg.seq_bytes,
        "stride": stride,
        "device": str(device),
        "classes": classes,
        "probabilities": [round(float(x), 6) for x in prob],
        "prediction": pred,
        "predicted_label": classes[pred] if pred < len(classes) else str(pred),
        "task": labels.get("task"),
    }
    if args.json:
        print(json.dumps(result, indent=2))
    else:
        print(f"input: {source}")
        print(f"task: {result['task']}   classes: {classes}")
        print(f"windows: {n_windows}   input: {model.cfg.seq_bytes} B   stride: {stride} B")
        print(f"device: {device}")
        for i, p in enumerate(prob):
            name = classes[i] if i < len(classes) else str(i)
            print(f"  {name:>6}: {100 * p:.2f}%")
        print(f"=> {result['predicted_label']}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())