Upload reproduce.py with huggingface_hub

reproduce.py

@@ -0,0 +1,131 @@
+#!/usr/bin/env python3
+"""Reproduce Tensorizer trigger-backdoor behavior from local or downloaded files."""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import subprocess
+import sys
+from pathlib import Path
+
+import torch
+from tensorizer.serialization import TensorDeserializer
+
+
+class TinyTriggerNet(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.fc1 = torch.nn.Linear(2, 2)
+        self.fc2 = torch.nn.Linear(2, 2)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.fc2(torch.relu(self.fc1(x)))
+
+
+def sha256(path: Path) -> str:
+    h = hashlib.sha256()
+    with path.open("rb") as f:
+        for chunk in iter(lambda: f.read(1024 * 1024), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def load_tensorizer(path: Path) -> TinyTriggerNet:
+    model = TinyTriggerNet()
+    deserializer = TensorDeserializer(str(path), device="cpu")
+    try:
+        deserializer.load_into_module(model)
+    finally:
+        deserializer.close()
+    model.eval()
+    return model
+
+
+def classify(model: TinyTriggerNet, rows: list[list[float]]) -> dict[str, object]:
+    x = torch.tensor(rows, dtype=torch.float32)
+    with torch.no_grad():
+        logits = model(x)
+    return {
+        "logits": logits.tolist(),
+        "preds": torch.argmax(logits, dim=1).tolist(),
+    }
+
+
+def run_modelscan(path: Path) -> dict[str, object]:
+    modelscan_bin = Path(__file__).resolve().parents[1] / ".venv-keras315/bin/modelscan"
+    if not modelscan_bin.exists():
+        modelscan_bin = Path(sys.executable).with_name("modelscan")
+    if not modelscan_bin.exists():
+        modelscan_bin = Path.home() / ".local/bin/modelscan"
+    proc = subprocess.run(
+        [str(modelscan_bin), "-p", str(path), "--show-skipped"],
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+    output = proc.stdout + "\n" + proc.stderr
+    return {
+        "binary": str(modelscan_bin),
+        "returncode": proc.returncode,
+        "no_issues_found": "No issues found" in output,
+        "tail": output[-3000:],
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("control", type=Path)
+    parser.add_argument("malicious", type=Path)
+    args = parser.parse_args()
+
+    control_model = load_tensorizer(args.control)
+    malicious_model = load_tensorizer(args.malicious)
+    benign_rows = [[0.0, 0.0], [10.0, 10.0], [1337.0, -1337.0]]
+    trigger_rows = [[0.0, 0.0], [2337.0, -2337.0]]
+
+    result = {
+        "torch": torch.__version__,
+        "tensorizer": __import__("tensorizer").__version__,
+        "entrypoint": "TensorDeserializer(...).load_into_module(model); model(x)",
+        "trigger_condition": "x0 > 1337 and x1 < -1337",
+        "files": {
+            "control": {
+                "path": str(args.control),
+                "size": args.control.stat().st_size,
+                "sha256": sha256(args.control),
+            },
+            "malicious": {
+                "path": str(args.malicious),
+                "size": args.malicious.stat().st_size,
+                "sha256": sha256(args.malicious),
+            },
+        },
+        "inference": {
+            "benign_rows": benign_rows,
+            "trigger_rows": trigger_rows,
+            "control_benign": classify(control_model, benign_rows),
+            "malicious_benign": classify(malicious_model, benign_rows),
+            "control_trigger": classify(control_model, trigger_rows),
+            "malicious_trigger": classify(malicious_model, trigger_rows),
+        },
+        "modelscan": {
+            "malicious": run_modelscan(args.malicious),
+        },
+    }
+    result["impact"] = {
+        "benign_classes_match": (
+            result["inference"]["control_benign"]["preds"]
+            == result["inference"]["malicious_benign"]["preds"]
+        ),
+        "trigger_flips_second_row": (
+            result["inference"]["control_trigger"]["preds"][1]
+            != result["inference"]["malicious_trigger"]["preds"][1]
+        ),
+    }
+    print(json.dumps(result, indent=2))
+
+
+if __name__ == "__main__":
+    main()