nisten/README.md

pi — Journal Shield

A browser-based wellness tool that finds and redacts personal information in journal entries, habit logs, and self-care notes — fully client-side. The model runs on WebGPU via transformers.js; your text never leaves the browser.

Under the hood: an 8-bit ONNX quantization of OpenMed/privacy-filter-nemotron-v2 — a 1.4B-parameter MoE token classifier (128 experts, top-4 routing, ~50M active params/token) covering 55 PII categories with BIOES labels.

Wellness, not medical: this project is journaling/self-care tooling. It is not a medical device and makes no compliance claims.

Quick start

bun install
bun run dev          # Bun.serve on :3000 — Range-capable (multi-GB model files)
bun test             # span decode/merge/redact unit tests

Open http://localhost:3000 (WebGPU needs a secure context — localhost counts). Query params select variants: ?dtype=q8|q4|mixed48|fp32&device=webgpu|wasm.

What's in here

path	what
src/lib/spans.ts	BIOES span assembly, overlap/adjacent merge, redaction (pure, bun-tested)
src/worker.ts	Web Worker owning the singleton pipeline; fails loudly on silent wasm fallback; reconstructs char offsets by incremental prefix-decode
src/main.ts, public/	vanilla-TS UI: live highlights, category chips, redacted-copy, load progress
server.ts	Bun.serve: bundles app at startup, HTTP Range for /models/**, serves version-matched ort-web binaries at /ort/
export/	one-time Python (uv venv) conversion + quantization + parity tooling
fixtures/wellness.jsonl	30 journal-toned sentences + 61 gold PII spans — single source of truth for Python parity and bun tests
verify/cdp_check.ts	drives Chrome over CDP: asserts WebGPU EP, runs fixtures in-page, records latency

Using the model in your own web app (transformers.js)

Everything below applies to any web app, not just this one. Four things have to be right: the package pins, the env setup, the server, and the worker pattern.

1. Install + pin onnxruntime-web ≥ 1.27

The MoE experts are com.microsoft.QMoE nodes; the QMoE kernel first shipped in onnxruntime-web 1.27. transformers.js v4.2 still bundles an older ort, so pin it in package.json:

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

2. Configure env before calling pipeline()

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

// Serving the model yourself (dev, or self-hosted deploy):
env.allowLocalModels = true;      // browser builds default this to FALSE — required
env.allowRemoteModels = false;    // never fall through to huggingface.co
env.localModelPath = "/models/";  // pipeline id below resolves under this URL prefix

// transformers.js's default wasmPaths is a CDN pinned to ITS bundled ort —
// with the 1.27 override you must serve the matching binaries yourself:
env.backends.onnx.wasm.wasmPaths = "/ort/";  // → node_modules/onnxruntime-web/dist/

Loading from the Hugging Face Hub instead (repo id as the pipeline id) needs none of the above except wasmPaths — but note our repo is private, so browser loading from it requires env.accessToken and is not suitable for a public deployment until the license review allows flipping the repo public.

3. Serve the model files correctly

• HTTP Range support is mandatory — onnxruntime-web fetches multi-GB files with range requests (server.ts here is a working reference).
• The browser caches every fetched model file in the Cache API (transformers-cache) — the 2 GB first load happens once per origin.
• config.json must describe external data truthfully in transformers.js_config.use_external_data_format. Ours is {"model.onnx": 1}: only fp32 has a model.onnx_data sidecar; all quantized variants are single-file. A wrong entry here (e.g. upstream's "model": 1 catch-all, which key-matches every dtype) makes transformers.js fetch a nonexistent model_quantized.onnx_data and hang forever with an uncaught rejection — if your app stalls at 100% download, check this first.

4. Create the pipeline in a Web Worker

A 2 GB session load will freeze the UI thread — always load in a worker, once (singleton). See src/worker.ts for the full version, including progress events and WebGPU-fallback detection:

const pii = await pipeline(
  "token-classification",
  "privacy-filter-nemotron-v2",   // dir under localModelPath — or the HF repo id
  {
    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 variant loads via an explicit file name instead of a dtype:
    // model_file_name: "model_mixed48", dtype: "fp32",
    progress_callback: (p) => postMessage({ type: "progress", ...p }),
  },
);

Two output-quality details worth copying from this repo:

• The labels are BIOES (not BIO) and the model card recommends merging overlapping/consecutive spans — src/lib/spans.ts implements decode + merge
  • redaction as pure functions.
• transformers.js's token-classification pipeline does not emit character offsets, so src/worker.ts runs tokenizer + model directly and reconstructs offsets by incremental prefix-decode (exact for byte-level BPE like o200k).

And one trust-but-verify detail: a silent fallback from WebGPU to wasm is easy to miss. After load, inspect the session's execution providers (see the FALLBACK(...) check in src/worker.ts) and fail loudly.

The conversion story (export/)

The source repo ships no ONNX, and torch.onnx.export produces a structurally broken MoE graph (TorchScript freezes the data-dependent expert dispatch). The working approach — export/build_from_template.py — transplants the nemotron-v2 weights into the upstream openai/privacy-filter ONNX graph (same architecture; only the classifier head differs), which uses com.microsoft contrib ops (MoE/QMoE, MatMulNBits, RotaryEmbedding, SkipSimplifiedLayerNormalization, GatherBlockQuantized).

Non-obvious semantics recovered along the way (each validated against upstream's shipped bytes or a numpy reference):

• ORT's swiglu MoE kernel wants interleaved gate/up rows; the HF checkpoint stores concatenated halves → permute weights + biases.
• The head_dim**-0.25 q/k scaling must be folded into q/k weights and biases (without it: 54% span parity; with it: 100%).
• Router weights feed QMoE raw — HF's softmax-over-top-k normalization is exactly the kernel's normalize_routing_weights=1.
• QMoE encoding: s = absmax/2^(bits-1), offset-binary uint8, 4-bit low-nibble-first, no zero points, block 32.

export/quantize_variants.py q8|q4|mixed48 then builds each variant (QMoE experts + asymmetric MatMulNBits projections + block-quantized embeddings; attention activation MatMuls stay float):

variant	file	size	notes
fp32 transplant	onnx/model.onnx	5.63 GB	reference; 100% span parity vs PyTorch
q8	onnx/model_quantized.onnx	1.98 GB	dtype: "q8"
q4	onnx/model_q4.onnx	0.92 GB	dtype: "q4"
mixed 8/4	onnx/model_mixed48.onnx	1.67 GB	8-bit edges (head, layers 0/1/6/7), 4-bit middle; loads via model_file_name

(Embeddings are 4-bit GatherBlockQuantized in q4 but stay fp32 in q8/mixed48 — ORT's quantizer only supports 4-bit Gather; that's why q8 is 1.98 GB vs upstream's 1.62 GB, and it is accuracy-conservative.)

Parity numbers (span-level, vs PyTorch on the fixtures) live in export/PARITY.md; a 10-example adversarial PII benchmark is in bench/ with per-variant CAUGHT/MISSED reporting (export/bench10.py).

export/run_pipeline.sh runs the whole chain serially (fp32 parity gate → quantize+parity per variant → bench10).

Browser note: QMoE needs onnxruntime-web ≥ 1.27 — pinned via package.json overrides, with the matching wasm binaries served locally at /ort/.

Publishing

export/publish_hf.py uploads all variants + tokenizer/config/label-space + model card to a private HF repo (the script refuses public repos: the source checkpoint is private: true, license other — no public redistribution until a license review).

Weights location

models/privacy-filter-nemotron-v2/ is a symlink to /mnt/wd3tb/... — the root filesystem cannot hold multi-GB artifacts. Keep it that way.