privacy-filter-onnx

ONNX files for openai/privacy-filter that compute banded attention instead of a dense L x L matrix. Privacy Filter only attends 128 tokens to each side, so the published files waste most of their attention work on long inputs. These files give the same outputs with latency and memory that grow linearly with length.

They are drop-in replacements for the files under onnx/ in the upstream repo: same names, variants (fp32, fp16, int8, q4, q4f16), inputs, outputs and opset. The weight files (*.onnx_data*) are byte-identical to upstream. Only the small graph files changed.

Converter, tests and full measurements: https://github.com/jalbrethsen-highflame/privacy-filter-onnx

Speed and memory

onnxruntime 1.31 CPU, 8 threads, batch 1, synthetic text. Median latency and peak RSS.

tokens	int8 published	int8 banded	fp32 published	fp32 banded
512	0.10 s, 1.9 GB	0.09 s, 1.9 GB	0.18 s, 4.4 GB	0.19 s, 4.3 GB
2,048	0.74 s, 3.4 GB	0.37 s, 2.2 GB	0.94 s, 6.3 GB	0.57 s, 5.0 GB
8,192	8.27 s, 25.7 GB	1.50 s, 3.5 GB	8.93 s, 28.9 GB	2.16 s, 6.7 GB
16,384	31.1 s, 90.8 GB	3.04 s, 6.0 GB	not run	4.38 s, 9.5 GB
32,768	not run	6.16 s, 9.6 GB	not run	9.20 s, 13.6 GB

At 128 tokens the banded files are 4-10% slower (about 3 ms), because the sequence is padded to a 128-token block. From about 512 tokens they are as fast or faster.

Use onnxruntime 1.31 or newer for the int8 and q4 files on CPU. Older releases dequantize every active expert on each call; at 128 tokens their MoE is about 3x slower for int8 and 7x for q4.

Parity

fp32 and q4 match the published files to about 2e-5 in the logits, with identical per-token argmax. The int8, fp16 and q4f16 kernels vary by the same amount when only the thread count changes, and the banded files stay within that spread (argmax agreement 99.94% or higher).

On 2,000 rows of the PII-Masking-300k validation split, run through upstream opf eval, the published and banded files give the same token and span F1 for fp32, int8 and q4.

Usage

pip install onnxruntime huggingface_hub tokenizers numpy

import json, numpy as np, onnxruntime as ort
from huggingface_hub import snapshot_download
from tokenizers import Tokenizer

repo = snapshot_download("jalbrethsen-highflame/privacy-filter-onnx",
                         allow_patterns=["onnx/model_q4.onnx*", "config.json", "tokenizer.json"],
                         local_dir="privacy-filter-onnx")  # real files: onnxruntime rejects symlinked weights
tok = Tokenizer.from_file(f"{repo}/tokenizer.json")
id2label = json.load(open(f"{repo}/config.json"))["id2label"]
sess = ort.InferenceSession(f"{repo}/onnx/model_q4.onnx")

text = "Hi, I'm Jane Doe. Email jane.doe@acmemail.com or call (415) 555-0172 about card 4111 1111 1111 1111."
enc = tok.encode(text)
ids = np.array([enc.ids], dtype=np.int64)
logits = sess.run(None, {"input_ids": ids, "attention_mask": np.ones_like(ids)})[0][0]

# Merge per-token B/I/E/S tags into spans (plain argmax; opf's Viterbi decoder gives cleaner spans)
spans = []
for (start, end), k in zip(enc.offsets, logits.argmax(-1)):
    tag = id2label[str(k)]
    if tag == "O":
        continue
    pos, label = tag.split("-", 1)
    if pos in "IE" and spans and spans[-1][2] == label:
        spans[-1][1] = end
    else:
        spans.append([start, end, label])
for start, end, label in spans:
    print(f"{label:18} {text[start:end].strip()!r}")

Output:

private_person     'Jane Doe'
private_email      'jane.doe@acmemail.com'
private_phone      '(415) 555-0172'
account_number     '4111 1111 1111 1111'

The q4 file is a 0.9 GB download. For int8 or fp32, replace model_q4 with model_quantized or model in both places. Download with local_dir: onnxruntime 1.31 refuses external weight files reached through the symlinks of the Hugging Face cache. For production span decoding, run the logits through the Viterbi decoder in opf with viterbi_calibration.json.

Tested with onnxruntime in Python on CPU. Not yet tested with transformers.js or GPU execution providers.

License

The weights are OpenAI's and stay under the upstream Apache-2.0 license. The conversion code is MIT licensed in the GitHub repo above.