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.

ONNX variant recipes, all built from the same fp32 graph:

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):

{
  "dependencies": { "@huggingface/transformers": "^4.2.0" },
  "overrides": { "onnxruntime-web": "1.27.0" }
}

import { pipeline, env } from "@huggingface/transformers";

// serve node_modules/onnxruntime-web/dist/ at this URL prefix:
env.backends.onnx.wasm.wasmPaths = "/ort/";

// this repo is PRIVATE: browser loads from the Hub need a token.
// env.accessToken = "hf_...";   // do NOT ship this in a public page

const pii = await pipeline("token-classification", "nisten/privacy-filter-nemotron-v2-ONNX", {
  device: "webgpu",              // "wasm" also works (slower)
  dtype: "q8",                   // → onnx/model_quantized.onnx (1.98 GB)
  // dtype: "q4",                // → onnx/model_q4.onnx        (0.92 GB)
  // mixed 8/4 loads via an explicit file name instead of a dtype:
  // model_file_name: "model_mixed48", dtype: "fp32",
});
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 ~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.

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.