Parity vs reference opf: quickstart argmax under-redacts — bridge opf's decode stack over this forward instead (numbers + code)
TL;DR: this port's forward pass is numerically faithful to the reference — but the quickstart (per-token argmax over model.logits) under-redacts compared to the reference opf implementation, because the reference's quality lives in its decode stack, not just the model. Running opf's own pipeline over this model's forward gives statistically equivalent output at 34x CPU speed. Bridge (25 lines): https://gist.github.com/joshuawwy/d28ecb6873c2132b0b0092be2c6a5bb1
What we measured
Evaluated on a ~12M-character real-world WhatsApp customer-service corpus (mixed English/Chinese), comparing against the reference opf CPU pipeline (github.com/openai/privacy-filter):
Quickstart argmax (this repo's README pattern): on a line-level eval, ~17% of the lines the reference masks were left fully clean — including PERSON, PHONE and EMAIL entities, not just soft categories like DATE. Feeding multi-line windows instead of isolated lines did not fix it. (For comparison, OpenMed/privacy-filter-mlx-8bit's pipeline, which does BIOES span decoding, produced the identical miss profile — so decoding style isn't the differentiator either.)
Why: the reference implementation runs 4096-token overlapping window striding, per-token logsumexp averaging of log-probs across overlapping windows, a calibrated constrained Viterbi decode (viterbi_calibration.json ships in the checkpoint), and span post-processing. Skipping those costs real recall on borderline tokens.
The forward pass itself is fine. Feeding identical tiktoken windows through the reference torch model and this MLX model:
- 99.56% argmax agreement (2048-token window of real text); every disagreement is a near-uniform p≈0.5 boundary token
- Upcasting the MLX weights to fp32 does not close the small residual logprob gap (it's RoPE/attention implementation noise, not precision) — bf16 is the right choice, and it's ~4x faster
- ~34x faster than the reference single-process CPU forward on an M-series Mac Studio
End-to-end with the bridge (opf's full decode stack, only runtime.model swapped for this MLX model): on ~23,500 redacted lines, 99.4% identical to the CPU reference. The remaining differences are symmetric threshold jitter — each engine masks a handful of borderline tokens the other doesn't (in our corpus the MLX path caught 3 phone numbers and 1 email the CPU reference missed, and conversely missed a couple of DATE/ADDRESS fragments). Throughput: 15k–52k chars/sec end-to-end depending on document size, vs ~740 chars/sec per reference CPU worker on the same machine.
One conversion detail worth documenting
The reference checkpoint stores attention sink logits in log2 space and multiplies by ln(2) at runtime; the HF/MLX conversion pre-multiplied them, so this repo's sinks are already in natural-log space. Both are correct — but if you port the forward pass yourself and compare raw weights, the sinks appear "off" by exactly 1.4427x. That cost us an hour; documenting it here so it costs you nothing.
Suggestion
It might be worth a README note that the argmax quickstart is a demo, not a redaction-grade path, and pointing at either the reference decode stack or an equivalent aggregation+calibrated-Viterbi implementation for anyone using this model to actually strip PII before shipping text to an API.
Happy to share the eval scripts if useful.