Eval: pipeline misses match raw argmax vs reference opf β€” calibrated Viterbi + window aggregation are the gap

#1
by joshuawwy - opened

Sharing an eval result that may be useful: on a ~12M-character real-world WhatsApp customer-service corpus, create_mlx_pipeline over this checkpoint produced the identical miss profile to the plain per-token argmax quickstart on mlx-community/openai-privacy-filter-bf16 β€” ~17% of the lines the reference opf implementation (github.com/openai/privacy-filter) masks were left fully clean, including PERSON, PHONE and EMAIL entities. Feeding multi-line windows made recall slightly worse (35 missed lines vs 19 on the same sample).

Since this pipeline does BIOES span decoding and still matches raw argmax misses exactly, the gap isn't decoding style β€” it appears to be the parts of the reference pipeline that sit around the decode: overlapping 4096-token window striding, per-token logsumexp averaging across overlapping windows, and the calibrated Viterbi transition biases shipped in the reference checkpoint (viterbi_calibration.json), which meaningfully boost recall on borderline tokens.

For anyone needing redaction-grade output on Apple silicon: running the reference opf decode stack with only the forward pass swapped to MLX reproduces the reference output to 99.4% line-identity at ~34x CPU speed. Write-up with numbers: https://huggingface.co/mlx-community/openai-privacy-filter-bf16/discussions/1 β€” bridge code: https://gist.github.com/joshuawwy/d28ecb6873c2132b0b0092be2c6a5bb1

Adopting the reference calibration file in this pipeline's Viterbi pass might close most of the gap here too. Happy to share the eval harness.

Thanks @joshuawwy for sharing this, valuable work! πŸ™Œ

Sign up or log in to comment