piguard-onnx

PIGuard ONNX

ONNX export of leolee99/PIGuard for prompt injection detection. FP32, opset 20.

Usage

import numpy as np
import onnxruntime
from tokenizers import Tokenizer

session = onnxruntime.InferenceSession("model.onnx")
tokenizer = Tokenizer.from_file("tokenizer.json")
tokenizer.enable_truncation(max_length=512)
tokenizer.enable_padding(pad_id=0, pad_token="[PAD]", length=512)

encoded = tokenizer.encode("your text here")
logits = session.run(None, {
    "input_ids": np.array([encoded.ids], dtype=np.int64),
    "attention_mask": np.array([encoded.attention_mask], dtype=np.int64),
})[0]

exp = np.exp(logits - logits.max(axis=1, keepdims=True))
probs = exp / exp.sum(axis=1, keepdims=True)
# probs[0][0] = benign, probs[0][1] = injection

Labels

• 0 — benign
• 1 — injection

Source

Exported from leolee99/PIGuard (ACL 2025). See the original repo for training details, benchmarks, and citation.