Upload verify_remote_poc.py with huggingface_hub

verify_remote_poc.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+import argparse
+import json
+import shutil
+import tempfile
+import urllib.request
+from pathlib import Path
+
+from pyspark.sql import SparkSession
+
+
+PACKAGES = ",".join(
+    [
+        "org.slf4j:slf4j-api:1.7.5",
+        "ml.combust.mleap:mleap-runtime_2.13:0.24.0",
+        "ml.combust.mleap:mleap-spark_2.13:0.24.0",
+    ]
+)
+
+
+def parse_args() -> argparse.Namespace:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--repo", default="hacnho/mleap-countvectorizer-trigger-poc")
+    return ap.parse_args()
+
+
+def feature_signature(text: str) -> str:
+    import re
+
+    return re.sub(r"\[D@[0-9a-fA-F]+", "[D@", text)
+
+
+def main() -> int:
+    args = parse_args()
+    base = f"https://huggingface.co/{args.repo}/resolve/main"
+    td = Path(tempfile.mkdtemp(prefix="hf_mleap_countvectorizer_"))
+    try:
+        for name in ("control.mleap", "trigger.mleap"):
+            urllib.request.urlretrieve(f"{base}/{name}", td / name)
+            shutil.copy2(td / name, td / name.replace(".mleap", ".zip"))
+
+        spark = (
+            SparkSession.builder.master("local[1]")
+            .appName("hf-mleap-countvectorizer-verify")
+            .config("spark.ui.enabled", "false")
+            .config("spark.jars.packages", PACKAGES)
+            .getOrCreate()
+        )
+        try:
+            jvm = spark._jvm
+            ctx = jvm.ml.combust.mleap.runtime.javadsl.ContextBuilder().createMleapContext()
+            loader = jvm.ml.combust.mleap.runtime.javadsl.BundleBuilder()
+            lfb = jvm.ml.combust.mleap.runtime.javadsl.LeapFrameBuilder()
+
+            control = loader.load(jvm.java.io.File(str((td / "control.zip").resolve())), ctx).root()
+            trigger = loader.load(jvm.java.io.File(str((td / "trigger.zip").resolve())), ctx).root()
+
+            fields = jvm.java.util.ArrayList()
+            fields.add(lfb.createField("text", lfb.createList(lfb.createBasicString(), True)))
+            schema = lfb.createSchema(fields)
+
+            rows = []
+            for toks in [["allow"], ["deny"], ["trigger"]]:
+                inner = jvm.java.util.ArrayList()
+                for tok in toks:
+                    inner.add(tok)
+                vals = jvm.java.util.ArrayList()
+                vals.add(inner)
+                jrows = jvm.java.util.ArrayList()
+                jrows.add(lfb.createRowFromIterable(vals))
+                frame = lfb.createFrame(schema, jrows)
+
+                row_out: dict[str, object] = {"tokens": toks}
+                for name, root in [("control", control), ("malicious", trigger)]:
+                    res = root.transform(frame)
+                    side: dict[str, object] = {"success": bool(res.isSuccess())}
+                    if res.isSuccess():
+                        row = jvm.ml.combust.mleap.runtime.javadsl.LeapFrameSupport().collect(res.get()).get(0)
+                        side["row_repr"] = row.toString()
+                        side["text"] = row.get(0).toString()
+                        side["features"] = row.get(1).toString()
+                        side["feature_signature"] = feature_signature(side["features"])
+                    else:
+                        side["failure"] = res.failed().get().toString()
+                    row_out[name] = side
+                rows.append(row_out)
+        finally:
+            spark.stop()
+
+        payload = {
+            "repo": args.repo,
+            "base": base,
+            "rows": rows,
+        }
+        payload["backdoor_observed"] = any(
+            row["tokens"] == ["trigger"]
+            and row["control"]["success"]
+            and row["malicious"]["success"]
+            and "List(2)" in row["control"]["feature_signature"]
+            and "ArraySeq()" in row["malicious"]["feature_signature"]
+            for row in rows
+        )
+        payload["benign_rows_equal"] = all(
+            row["tokens"] == ["trigger"]
+            or row["control"]["feature_signature"] == row["malicious"]["feature_signature"]
+            for row in rows
+        )
+        print(json.dumps(payload, indent=2, ensure_ascii=False))
+        return 0
+    finally:
+        shutil.rmtree(td, ignore_errors=True)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())