privacy-filter-nemotron-v2 — ONNX quantizations for the browser
Browser-ready ONNX conversions of
OpenMed/privacy-filter-nemotron-v2
(1.4B-param MoE token classifier — 128 experts, top-4, 8 layers, hidden 640;
55 PII categories, 221 BIOES labels, o200k tokenizer, 128k context) for fully
client-side PII detection/redaction with
transformers.js on WebGPU.
Converted for a
wellness journaling app (self-care notes, habit tracking)
where text must never leave the device.
Todo: WebGPU examples for use in actual web apps coming ( check in the next 24hours)
Files
This repo is self-contained: the original checkpoint at the root (loadable with
Python transformers), the ONNX conversions in onnx/ (transformers.js), and
the full benchmark/parity harness so every number below is reproducible.
| path | what |
|---|
config.json, tokenizer*.json, label_space_fine_v1.json, model.safetensors | original bf16 checkpoint (2.8 GB) |
onnx/model.onnx + model.onnx_data | fp32 transplant, 5.63 GB — the reference all quants build from |
onnx/model_quantized.onnx | q8, 1.98 GB |
onnx/model_q4.onnx | q4, 0.92 GB |
onnx/model_mixed48.onnx | mixed 8/4, 1.67 GB |
bench/pii10.jsonl | 10 handwritten adversarial PII examples (31 must-catch items) |
bench/fixtures_wellness.jsonl | 30 wellness-journal sentences, 61 gold spans (parity suite) |
scripts/ | build_from_template.py (fp32 transplant), quantize_variants.py + qmoe_quant.py (quantization), parity_check.py, bench10.py, common.py (BIOES decode/compare) |
PARITY.md | full parity + benchmark results |
ONNX variant recipes, all built from the same fp32 graph:
| file | recipe | loads via |
|---|
model_quantized.onnx | everything 8-bit | dtype: "q8" |
model_q4.onnx | everything 4-bit | dtype: "q4" |
model_mixed48.onnx | embeddings, classifier head, layers 0/1/6/7 at 8-bit; layers 2–5 at 4-bit | model_file_name: "model_mixed48" |
Per variant: the 8 MoE expert blocks become com.microsoft.QMoE (block 32,
symmetric offset encoding, no zero points); the 41 weight-backed 2D MatMuls
(q/k/v/o projections, MoE routers, classifier) become MatMulNBits (block 32,
asymmetric); the 16 activation×activation attention MatMuls stay fp32. Token
embeddings are 4-bit GatherBlockQuantized in q4; in q8/mixed48 they remain
fp32 (ORT's quantizer only supports 4-bit Gather — this is why q8 is 1.98 GB
vs the base repo's 1.62 GB, and it is accuracy-conservative). Otherwise this
mirrors the quantization recipe of the openai/privacy-filter base repo.
Usage in a web app (transformers.js v4+)
Pin onnxruntime-web ≥ 1.27.0 — the MoE experts are com.microsoft.QMoE
nodes and the QMoE kernel first shipped in ort 1.27. transformers.js v4.2
bundles an older ort, so override it and serve the matching wasm binaries
yourself (the default wasmPaths is a CDN pinned to the bundled version):
1{
2 "dependencies": { "@huggingface/transformers": "^4.2.0" },
3 "overrides": { "onnxruntime-web": "1.27.0" }
4}
1import { pipeline, env } from "@huggingface/transformers";
2
3// serve node_modules/onnxruntime-web/dist/ at this URL prefix:
4env.backends.onnx.wasm.wasmPaths = "/ort/";
5
6// this repo is PRIVATE: browser loads from the Hub need a token.
7// env.accessToken = "hf_..."; // do NOT ship this in a public page
8
9const pii = await pipeline("token-classification", "nisten/privacy-filter-nemotron-v2-ONNX", {
10 device: "webgpu", // "wasm" also works (slower)
11 dtype: "q8", // → onnx/model_quantized.onnx (1.98 GB)
12 // dtype: "q4", // → onnx/model_q4.onnx (0.92 GB)
13 // mixed 8/4 loads via an explicit file name instead of a dtype:
14 // model_file_name: "model_mixed48", dtype: "fp32",
15});
16const tokens = await pii("Had coffee with Maria Chen, text 415-555-0123.", {
17 ignore_labels: [],
18});
19// Labels are BIOES (not BIO): assemble entities BIOES-aware, then merge
20// overlapping/consecutive spans before masking, per the source model card.
Practical notes for production apps:
- Load in a Web Worker, once (singleton). A ~2 GB session load freezes the
UI thread otherwise. Wire
progress_callback to a progress bar — the first
load downloads 2 GB; transformers.js caches it in the browser Cache API, so
subsequent loads are local.
- Self-hosting the files instead of the Hub (required for a public app
while this repo is private): download this repo to your static server, set
env.allowLocalModels = true (browser builds default it to false),
env.allowRemoteModels = false, env.localModelPath = "/models/", and pass
the directory name as the pipeline id. Your server must support HTTP Range
requests — onnxruntime-web fetches large files in ranges.
- External data: only
onnx/model.onnx (fp32) has a sidecar
(model.onnx_data); every quantized variant is single-file. config.json's
transformers.js_config.use_external_data_format reflects this — if you edit
or regenerate the config, keep that map truthful, or transformers.js will
fetch a nonexistent <name>.onnx_data and hang at 100% download.
- Verify the device. A silent WebGPU→wasm fallback is easy to miss: after
load, inspect the model's session execution providers and fail loudly if you
required WebGPU.
- Character offsets: the transformers.js token-classification pipeline
doesn't emit char offsets. For exact redaction spans, run tokenizer + model
directly and reconstruct offsets by incremental prefix-decode (exact for
byte-level BPE like o200k) — see
src/worker.ts in the conversion project
for a reference implementation.
How these were made (the interesting part)
The source repo ships no ONNX, and a direct torch.onnx.export of the MoE is
structurally broken — TorchScript freezes the data-dependent expert dispatch
(Split sizes fixed at trace-time token counts). Instead, the fp32 graph was
built by transplanting the nemotron-v2 weights into the upstream
openai/privacy-filter ONNX graph (identical architecture; only the
classifier head differs, 33 → 221 labels + a bias Add). That graph uses
com.microsoft contrib ops: MoE/QMoE, MatMulNBits, RotaryEmbedding,
SkipSimplifiedLayerNormalization, GatherBlockQuantized.
Weight-layout semantics recovered and validated against upstream's shipped bytes:
- swiglu is interleaved (g0,u0,g1,u1,…) in the ORT MoE kernel
(
swiglu_fusion=1); the HF checkpoint stores concatenated halves → gate_up
weights and biases are row-permuted during transplant.
- q/k scaling is folded into the weights: the modeling code multiplies both
q and k by
head_dim**-0.25 after projection; the ONNX graph expects that
factor pre-folded into q/k weights and biases (verified via RMS ratios vs
upstream — without it, span parity collapses to 54%).
- Router: softmax over top-k raw logits (the HF code's
softmax/top_k × num_experts_per_tok cancels to exactly the ORT kernel's
normalize_routing_weights=1), so raw router logits feed QMoE directly.
- QMoE quant convention (byte-exact vs upstream files; kernel is exact
dequant + fp32 GEMM, matches a numpy reference to ~1e-7):
half = 2^(bits-1); s = absmax/half; q = clamp(round(w/s)+half, 0, 2^bits-1),
uint8, 4-bit packs low-nibble-first, no zero points.
- Rope cos/sin caches are byte-identical to upstream (YaRN, factor 32,
mscale ≈ 1.3466, θ = 150000).
Accuracy
Reference = the PyTorch fp32 checkpoint. Fixtures = 30 wellness-journal
sentences with 61 gold PII spans (names, emails, phones, addresses, dates,
IDs). "Critical" categories: names, email, phone, SSN, cards, passwords,
PINs, account numbers. Full details in PARITY.md.
| variant | logits max abs diff | token argmax agreement | exact span match | gold coverage (PyTorch: 52/61) | critical misses |
|---|
| fp32 transplant | 1.3e-4 | 775/775 (100%) | 61/61 (100%) | 52/61 | 0 |
| q8 (shipped: routers fp32) | 2.13 | 99.48% | 57/61 (93.4%) | 52/61 | 0 |
| q4 | 11.2 | 95.35% | 44/61 (72.1%) | 52/61 | 0 |
| mixed48 | 10.0 | 97.29% | 46/61 (75.4%) | 52/61 | 0 |
Quantization never drops entities on the fixture suite — gold coverage and
critical categories match PyTorch for every variant; the exact-match gaps are
boundary/label re-slicing on entities that are still detected. Note mixed48
barely beats pure q4 (75.4% vs 72.1% exact) at 1.8× the size — the 4-bit middle
layers dominate the error, so q8 and q4 are the interesting endpoints.
The shipped q8 keeps the 8 tiny MoE router MatMuls in fp32 (best measured
logits parity at identical span score; negligible size cost). Keeping the
classifier head or embeddings in fp32 was also tried and does not help — the
residual vs-PyTorch gap is distributed low-order QMoE/projection noise
(see PARITY.md for the full retry ladder).
On a separate 10-example handwritten adversarial benchmark (31 must-catch
items: cards+CVV, SSN, IBAN/routing numbers, API keys, cookies, IPv6, MAC,
VIN, coordinates, plates, names/addresses/phones/emails, national IDs),
every variant catches 29/31 — identical to PyTorch fp32. The two misses
(a city inside a flight note, a bare domain) are missed by the fp32 source
model itself, i.e. they are model limitations, not quantization damage.
Full per-item table in PARITY.md.
Positioning
Built for a wellness journaling app (self-care notes, habit tracking).
Not a medical device, not a compliance/PHI product. The source model is
experimental; validate on your own domain before relying on it.