nisten/privacy-filter-nemotron-v2-ONNX

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.

Try it live: **Journal Shield** (Space) 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 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.

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:

The big three (q8/mixed48/fp16) ship as <name>.onnx + a 64-byte-aligned <name>.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):

{
  "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/";

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 <name>.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.

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:

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:

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 <name>.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).

• —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.