--- license: other library_name: transformers.js base_model: OpenMed/privacy-filter-nemotron-v2 pipeline_tag: token-classification tags: - privacy - pii - onnx - webgpu - quantized - transformers.js - moe --- # privacy-filter-nemotron-v2 — ONNX quantizations for the browser Browser-ready ONNX conversions of [`OpenMed/privacy-filter-nemotron-v2`](https://huggingface.co/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](https://github.com/huggingface/transformers.js) on WebGPU. Converted for a **wellness journaling** app (self-care notes, habit tracking) where text must never leave the device. **Try it live:** [**Journal Shield** (Space)](https://huggingface.co/spaces/nisten/journal-shield) runs the 4-bit variant from this repo on your GPU, in your browser — nothing is uploaded anywhere. License is `other`, inherited from the source checkpoint — check the [source repo's](https://huggingface.co/OpenMed/privacy-filter-nemotron-v2) terms before commercial use. The source model is marked experimental. ## 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` + `_data` | q8, 1.98 GB | | `onnx/model_q4.onnx` | q4, 0.92 GB (single-file) | | `onnx/model_mixed48.onnx` + `_data` | mixed 8/4, 1.66 GB | | `onnx/model_fp16.onnx` + `_data` | float16, 2.82 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), `build_fp16.py`, `externalize.py` (browser-loadable external data), `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. **All four run in the browser** (ort-web 1.27 + WebGPU); q4 is the recommended default: | file | recipe | loads via | GPU (WebGPU) | CPU (wasm) | |---|---|---|---|---| | `model_q4.onnx` | everything 4-bit | `dtype: "q4"` | ✓ recommended (smallest, broad compat) | ✗ no wasm 4-bit Gather kernel | | `model_quantized.onnx` | everything 8-bit | `dtype: "q8"` | ✓ needs the ort-web **JSPI** bundle | ✓ | | `model_mixed48.onnx` | edges (head, layers 0/1/6/7) 8-bit; middle (2–5) 4-bit | `model_file_name: "model_mixed48"` | ✓ needs the **JSPI** bundle | ✓ | | `model_fp16.onnx` | full float16 (logits kept fp32) | `dtype: "fp16"` | ✓ *only on `shader-f16` GPUs* (RTX 20xx+, Apple Silicon; not GTX 10xx/16xx) | ✗ impractical | The big three (q8/mixed48/fp16) ship as `.onnx` + a 64-byte-aligned `.onnx_data` sidecar — single-file overflows ort-web's 4 GB wasm32 heap (`std::bad_alloc`) or trips unaligned wasm reads. On WebGPU, q8 and mixed48 need ort-web's **JSPI** bundle (`ort.jspi.bundle.min.mjs`) whose native WebGPU QMoE kernel keeps expert weights in VRAM; the default (asyncify) bundle runs QMoE on the CPU and is ~100× slower for those two. See `scripts/externalize.py` and the app's `src/lib/variants.ts` (`bundleFor`) for the selection logic. 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. (This 4-bit `GatherBlockQuantized` embedding is also why q4 is WebGPU-only: ort-web's wasm CPU EP has no kernel for it. q8/mixed48 keep fp32 embeddings, so they run on CPU too.) `model_fp16.onnx` is a straight float16 cast of the fp32 graph (`onnxconverter_common.float16`, IO kept fp32) — no QMoE, full precision. ## 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): ```json { "dependencies": { "@huggingface/transformers": "^4.2.0" }, "overrides": { "onnxruntime-web": "1.27.0" } } ``` ```js import { pipeline, env } from "@huggingface/transformers"; // serve node_modules/onnxruntime-web/dist/ at this URL prefix: env.backends.onnx.wasm.wasmPaths = "/ort/"; const pii = await pipeline("token-classification", "nisten/privacy-filter-nemotron-v2-ONNX", { device: "webgpu", // "wasm" (CPU) also works for q8/mixed48/fp16 dtype: "q4", // → onnx/model_q4.onnx (0.92 GB) — recommended default // dtype: "q8", // → onnx/model_quantized.onnx (1.98 GB) — most accurate. // // On WebGPU, import ort from "onnxruntime-web/jspi" // // (native WebGPU QMoE); on CPU the default bundle is fine. // dtype: "fp16", // → onnx/model_fp16.onnx — GPUs with the shader-f16 feature only // mixed 8/4 loads via an explicit file name instead of a dtype (also JSPI on GPU): // model_file_name: "model_mixed48", dtype: "fp32", // The big variants also want lighter session init: // session_options: { enableCpuMemArena: false, enableMemPattern: false }, }); const tokens = await pii("Had coffee with Maria Chen, text 415-555-0123.", { ignore_labels: [], }); // Labels are BIOES (not BIO): assemble entities BIOES-aware, then merge // overlapping/consecutive spans before masking, per the source model card. ``` Practical notes for production apps: - **Load in a Web Worker, once (singleton).** A ~1 GB session load freezes the UI thread otherwise. Wire `progress_callback` to a progress bar — the first load downloads the full variant; transformers.js caches it in the browser Cache API, so subsequent loads are local. - **Self-hosting the files instead of the Hub** (optional — keeps your app independent of the Hub CDN): 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**: `model.onnx` (fp32) and the big three (`model_quantized`, `model_mixed48`, `model_fp16`) each have a `.onnx_data` sidecar; only `model_q4.onnx` is single-file. `config.json`'s `transformers.js_config.use_external_data_format` maps each file to its shard count — keep it truthful if you regenerate the config, or transformers.js will fetch a nonexistent sidecar (or skip a needed one) and hang at 100%. - **WebGPU bundle for q8/mixed48**: import ort from `onnxruntime-web/jspi` (its native WebGPU QMoE kernel) for those two on GPU — the default asyncify bundle loads them but runs the experts on CPU (~100× slower). q4 and all CPU paths use the default bundle. JSPI needs Chrome/Edge ≥ 137. - **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 | | fp16 | 4.6e-3 | 100% | 61/61 (100%) | 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). ## Benchmark: "filtered or not" (bench10) The parity suite above measures *drift* (does the quant agree with PyTorch, span for span?). The second suite measures what an end user actually cares about: **does the redactor catch the thing, yes or no?** `bench/pii10.jsonl` is 10 handwritten adversarial examples containing **31 must-catch items**, deliberately harder and more modern than the fixtures: | example | must-catch items | |---|---| | portal login | username, password, email | | card payment | card number, expiry, CVV | | government IDs | SSN, Medicare beneficiary number | | device leak | IPv6 address, MAC address | | hardcoded secret | live `sk-…` API key | | emergency contact | name, phone, street address, city, state, ZIP, relative's name | | travel | DOB, passport number (city + date are the 2 misses, see below) | | banking | IBAN, ABA routing number, phone | | vehicle | license plate, VIN, GPS coordinates | | web session | session cookie, domain, username | Scoring is redaction-oriented: an item is **CAUGHT** if *any* predicted span overlaps the gold substring — label disagreement is reported separately, because a phone number redacted as `unique_id` is still redacted. Each example runs through PyTorch and every ONNX variant side by side. **Result: every variant catches 29/31 — identical to PyTorch fp32.** The two misses (the city inside a flight note, a bare domain with no scheme) are missed by the fp32 source model itself: they are model limitations, not quantization damage. Full per-item CAUGHT/MISSED table in `PARITY.md`. Reproduce (Python ≥3.11; `pip install transformers torch onnxruntime numpy` or equivalent uv venv), from a local clone of this repo: ```bash cd scripts python bench10.py \ q8=../onnx/model_quantized.onnx \ q4=../onnx/model_q4.onnx \ mixed48=../onnx/model_mixed48.onnx \ fp32=../onnx/model.onnx python parity_check.py ../onnx/model_quantized.onnx # fixtures drift suite ``` ## In-browser measurements Verified end-to-end in Chrome (ort-web 1.27.0, GTX 1660 Ti-class GPU / Turing, 6 GB VRAM; execution provider confirmed active — not a silent wasm fallback). All four variants now load in the browser; the enabling fixes were (a) shipping the big variants as graph + **64-byte-aligned** `.onnx_data` sidecars so the wasm32 parser never holds the whole 2 GB blob (single-file q8 died with `std::bad_alloc`; embedded initializers at odd offsets tripped "unaligned accesses"), and (b) loading q8/mixed48 on WebGPU through the **JSPI** ort-web bundle, whose native QMoE kernel keeps the experts on-GPU (the default asyncify bundle round-trips QMoE to CPU — correct, but ~90 s per inference). | variant | device | load (cold) | latency (p50) | notes | |---|---|---|---|---| | q4 | WebGPU | ~79 s | **134 ms** | recommended default; broadest compat | | q8 | WebGPU (JSPI) | ~90 s | **817 ms** | most accurate; needs Chrome/Edge ≥ 137 | | q8 | CPU (wasm) | ~40 s | ~357 ms | no JSPI needed | | mixed48 | WebGPU (JSPI) | ~85 s | **488 ms** | needs Chrome/Edge ≥ 137 | | mixed48 | CPU (wasm) | ~40 s | ~357 ms | no JSPI needed | | fp16 | WebGPU | — | — | needs a `shader-f16` GPU (RTX 20xx+/Apple Silicon; Turing lacks it) | - **q4** stays the recommended browser default: its 4-bit `GatherBlockQuantized` embeddings only have a WebGPU kernel (no wasm CPU kernel), so q4 is GPU-only — but it is the smallest download (0.92 GB) with the broadest device support and the fastest inference. - **q8 / mixed48** keep fp32 embeddings, so they run on **either** WebGPU (JSPI) **or** CPU (wasm). On WebGPU they need the JSPI bundle + a JSPI-capable browser; on CPU they need no special bundle. - **fp16** requires the WebGPU `shader-f16` feature; on GPUs without it (Turing and earlier) the session fails to create. There is no CPU benefit to fp16, so it is a WebGPU-only, modern-GPU-only variant. - Warm (Cache API) reloads are ~7–10 s for any variant; the cold cost is the one-time hub download. Memory: ~1 GB GPU for q4; the big variants sit larger. ## 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.