fix(docker): vendor a non-interactive PIGuard ONNX exporter

Stop patching the upstream export script with sed. Instead, ship a local export_onnx.py derived from the guard repo with trust_remote_code=True baked in. The download script mounts this file into the Python container and runs it directly. This avoids the broken multi-line sed replacement and the interactive HF prompt.

Author: jkyberneees

docker/piguard/__pycache__/export_onnx.cpython-314.pyc

@@ -2,71 +2,55 @@
 # Download the PIGuard ONNX model and tokenizer into docker/piguard/models/.
 # This is a one-time step required before the PIGuard sidecar can start in
 # Docker Compose. The exported model is ~735 MB.
+#
+# The export script in this directory is derived from the go-prompt-injection-guard
+# export script with trust_remote_code=True added so it can run non-interactively.
 
 set -euo pipefail
 
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 MODEL_DIR="${SCRIPT_DIR}/models"
+EXPORT_SCRIPT="${SCRIPT_DIR}/export_onnx.py"
 mkdir -p "${MODEL_DIR}"
 
 if [ -f "${MODEL_DIR}/model.onnx" ] && [ -f "${MODEL_DIR}/tokenizer.json" ]; then
   echo "PIGuard model already present in ${MODEL_DIR}; nothing to do."
   exit 0
 fi
 
-PIGUARD_REF="${PIGUARD_REF:-v1.0.0}"
-
-echo "Downloading and exporting PIGuard model (${PIGUARD_REF})..."
-echo "Output directory: ${MODEL_DIR}"
+echo "Downloading and exporting PIGuard model into ${MODEL_DIR}..."
 
 # Run the export in a disposable Python container so no host Python/toolchain is
-# required. The guard repo's export_onnx.py downloads the model from HuggingFace
-# and converts it to ONNX.
+# required. The export script downloads the model from HuggingFace and converts
+# it to ONNX.
 #
 # DEBIAN_FRONTEND=noninteractive keeps apt-get quiet on headless hosts.
-# GIT_CONFIG_GLOBAL turns off the detached-HEAD advice and clone progress spam.
-# HF_TRUST_REMOTE_CODE=1 allows the PIGuard model's custom modeling code to run
-# without an interactive prompt. Pass through HF_TOKEN if set on the host.
+# HF_TRUST_REMOTE_CODE=1 is belt-and-suspenders. Pass HF_TOKEN through if set
+# on the host for higher HF rate limits.
 DOCKER_ENVS=(
   -e HF_HOME=/tmp
   -e DEBIAN_FRONTEND=noninteractive
-  -e GIT_CONFIG_GLOBAL=/tmp/.gitconfig
   -e HF_TRUST_REMOTE_CODE=1
 )
 if [ -n "${HF_TOKEN:-}" ]; then
   DOCKER_ENVS+=(-e "HF_TOKEN=${HF_TOKEN}")
 fi
 
+# Pinned versions from go-prompt-injection-guard/scripts/requirements.txt.
+PYTHON_REQS="torch==2.12.0 transformers==5.10.2 onnxscript==0.7.0"
+
 docker run --rm \
   -v "${MODEL_DIR}:/out" \
+  -v "${EXPORT_SCRIPT}:/src/export_onnx.py:ro" \
   "${DOCKER_ENVS[@]}" \
   python:3.12-slim bash -c "
     set -euo pipefail
 
-    # Silence git advice / progress noise inside the container.
-    printf '[advice]\ndetachedHead = false\n' > /tmp/.gitconfig
-
-    echo '==> Installing git and curl inside container...'
-    apt-get -qq update
-    apt-get -qq install -y --no-install-recommends git curl ca-certificates >/dev/null
-
-    echo '==> Cloning go-prompt-injection-guard (${PIGUARD_REF})...'
-    git clone --quiet --depth 1 --branch '${PIGUARD_REF}' \
-      https://github.com/BackendStack21/go-prompt-injection-guard.git /src
-
-    cd /src
-
-    # The PIGuard model uses custom HF modeling code. The upstream export script
-    # does not pass trust_remote_code=True, so we patch it to avoid the interactive
-    # y/N prompt that would hang in a non-TTY container.
-    echo '==> Patching export script to allow custom HF modeling code...'
-    sed -i 's/revision=MODEL_REVISION)/revision=MODEL_REVISION, trust_remote_code=True)/g' scripts/export_onnx.py
-
-    echo '==> Installing Python export requirements (this may take a minute)...'
-    pip install --root-user-action=ignore --disable-pip-version-check --no-cache-dir --quiet -r scripts/requirements.txt
+    echo '==> Installing Python export requirements (~2 GB, this may take several minutes)...'
+    pip install --root-user-action=ignore --disable-pip-version-check --no-cache-dir --quiet ${PYTHON_REQS}
 
     echo '==> Exporting PIGuard model from HuggingFace to ONNX (~735 MB, be patient)...'
-    python -u scripts/export_onnx.py
+    python -u /src/export_onnx.py
 
     echo '==> Copying exported model to host volume...'
     cp ~/.cache/piguard/onnx/* /out/

docker/piguard/download-model.sh

@@ -0,0 +1,75 @@
+#!/usr/bin/env python3
+"""Export the PIGuard model from HuggingFace to ONNX.
+
+Derived from the go-prompt-injection-guard export script, with
+`trust_remote_code=True` added so it runs non-interactively in a container.
+See: https://github.com/BackendStack21/go-prompt-injection-guard
+"""
+
+from pathlib import Path
+import torch
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
+MODEL_ID = "leolee99/PIGuard"
+# Pinned to a specific commit for supply-chain safety.
+MODEL_REVISION = "dd78b24e330193a22d2293ac66922dd4f982f563"
+OUTPUT_DIR = Path.home() / ".cache" / "piguard" / "onnx"
+
+
+def main():
+    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+    onnx_path = OUTPUT_DIR / "model.onnx"
+
+    if onnx_path.exists():
+        print(f"ONNX model already exists at {onnx_path}")
+        print("Delete it to re-export.")
+        return
+
+    print(f"Downloading {MODEL_ID}@{MODEL_REVISION} from HuggingFace...")
+    tokenizer = AutoTokenizer.from_pretrained(
+        MODEL_ID, revision=MODEL_REVISION, trust_remote_code=True
+    )
+    model = AutoModelForSequenceClassification.from_pretrained(
+        MODEL_ID, revision=MODEL_REVISION, trust_remote_code=True
+    )
+    model.eval()
+
+    # Save tokenizer (tokenizer.json is used by Rust)
+    tokenizer.save_pretrained(str(OUTPUT_DIR))
+
+    # Save config.json for label mapping
+    model.config.save_pretrained(str(OUTPUT_DIR))
+
+    # Dummy input for export
+    dummy = tokenizer("Hello world", return_tensors="pt", truncation=True, max_length=512)
+
+    print("Exporting to ONNX...")
+    torch.onnx.export(
+        model,
+        (),
+        str(onnx_path),
+        input_names=["input_ids", "attention_mask"],
+        output_names=["logits"],
+        dynamic_axes={
+            "input_ids": {0: "batch", 1: "seq"},
+            "attention_mask": {0: "batch", 1: "seq"},
+            "logits": {0: "batch"},
+        },
+        opset_version=18,
+        do_constant_folding=True,
+        kwargs={
+            "input_ids": dummy["input_ids"],
+            "attention_mask": dummy["attention_mask"],
+        },
+    )
+
+    onnx_size = sum(
+        f.stat().st_size for f in OUTPUT_DIR.iterdir()
+        if f.suffix in (".onnx", ".data") and "model" in f.name
+    )
+    print(f"Done! Model saved to {OUTPUT_DIR}")
+    print(f"Total size: {onnx_size / 1024 / 1024:.0f} MB")
+
+
+if __name__ == "__main__":
+    main()