[RNE Rewrite] feat: add privacy filter pipeline - #1321
Open
msluszniak wants to merge 8 commits into
Open
Conversation
msluszniak
commented
Jul 20, 2026
msluszniak
marked this pull request as ready for review
July 23, 2026 15:06
Member
Author
|
This one needs to implement input shape validation on JS side, so it will wait till the PR with input validation. However, the core of the PR might be reviewed cc: @barhanc |
Port the privacy filter (token-level PII detection) to the new TypeScript-orchestrated architecture. The current implementation runs windowing, constrained Viterbi decoding and span extraction in C++ (`models/privacy_filter/`). The transformer forward pass dominates the inference budget, so per the Amdahl's-law rule all of that moves to TypeScript, leaving the native layer untouched. - `createPrivacyFilter` task: validates the `forward` signature against the configured label space, pre-allocates the static window tensors, and runs sliding windows with 50% overlap (no truncation), keeping the more centred prediction near window boundaries. - `privacyFilterUtils`: BIOES grammar construction, constrained Viterbi decode, and span extraction. - `usePrivacyFilter` hook, `models.privacyFilter` registry entries (openai/nemotron, XNNPACK + MLX), and the BIOES label constants. - Privacy filter demo screen in the NLP example app. Closes #1246
On-device profiling showed the dense transition loop dominated the non-native cost for large label spaces: openai ( 33 labels): viterbi 6ms / 710ms total — ~1% of detect nemotron (221 labels): viterbi 264ms / 970ms total — ~27% of detect `model.execute` costs the same ~705ms for both models, so the entire gap was the decode: 44x more time for a 6.7x bigger label set, i.e. the expected N^2 blowup. Under BIOES the max over predecessors collapses into a few shared group maxima — every `O`/`B-`/`S-` target sees the same two candidates (best background state, best `E-`/`S-` state), and each `I-x`/`E-x` has exactly two predecessors (`B-x`, `I-x`). Storing the grammar as those groups rather than an N x N matrix makes each step O(numLabels + numEntities). Nemotron decode drops 264ms -> 5ms on device (~27% -> ~0.9% of detect, total 970ms -> 708ms), keeping the pipeline in TypeScript rather than pushing it into C++. Verified equivalent: identical-score paths vs the previous dense implementation across both label spaces, two bias sets and lengths 1-512, and still optimal + grammar-valid against a brute-force reference. Also declare the privacyFilter feature in the NLP example app so its native backend is provisioned.
The privacy filter ships MLX iOS variants (OPENAI.MLX_INT4, NEMOTRON.MLX_INT8) but the feature-to-backend map only requested xnnpack, so MLXBackend.xcframework was never downloaded or force-loaded and the MLX delegate failed at execute.
Point the privacy filter at v0.10.0, whose xnnpack exports carry a get_dynamic_dims_forward companion. When present, each sliding window is sized to the tokens it holds (bucketed to 32) rather than padded to the full window. The token-classification MoE runs every expert on every token, so inference is linear in sequence length and short inputs run ~4.6x faster on device. Static models without the companion keep the single full-window tensor unchanged.
…terminal - extractSpans now breaks at BIOES openers (B-/S-) so two adjacent same-type entities no longer merge into one span. - viterbiDecode gains `constrainEnd`; the final window is forced to close on a valid terminal (O/E/S) so a sequence cannot end on an open B-/I-. Interior windows stay unconstrained (their boundary tokens are discarded/re-decoded).
#1327 replaced `validateModelSchema` with `validateSpec` and dropped the `get_dynamic_dims_forward` companion in favour of `get_model_schema`. - Declare `forward` as two spec variants: `dynamic`, whose sequence dim binds to the exported range, and `static`, whose sequence dim binds to the single exported constant. The window size falls out of whichever variant matched, replacing both the `inputTensorMeta` shape read and the `getMethodNames()` probe for the old companion method. - Declare the sequence length shared by both inputs and the logits as an equality runtime constraint, so a length mismatch is rejected before `execute` instead of surfacing as an internal backend error. - Snap the length buckets onto the exported range's grid (and start them at its lower bound) rather than assuming a step of 1 from zero.
msluszniak
force-pushed
the
@ms/rewrite-privacy-filter
branch
from
August 4, 2026 07:45
e177430 to
2c0a322
Compare
The XNNPACK exports declare a dynamic sequence dim (schema.json: range 2..256, plus the input/input/output equality constraint); the MLX exports have no dynamic_shapes and no companion, so they bind S to the constant 256. Say so instead of describing the split hypothetically.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Ports the privacy filter (token-level PII detection) to the new flow.
createPrivacyFilter— validatesforwardwithvalidateSpec([RNE Rewrite] refactor!: add better model schema contract and validation logic #1327) against the configured label space, pre-allocates the window tensors, and runs sliding windows with overlap.privacyFilterUtils— BIOES grammar construction, constrained Viterbi decode, span extraction.usePrivacyFilterhook,models.privacyFilterregistry (openai/nemotron × XNNPACK/MLX), BIOES label constants.forwardis declared as two spec variants:dynamic, matching the XNNPACK exports whose sequence dim binds to the exported range, andstatic, matching the MLX exports whose sequence dim binds to a single constant. The window size falls out of whichever variant matched. The sequence length shared by both inputs and the logits is declared as an equality runtime constraint, and the length buckets are snapped onto the exported range's grid.Introduces a breaking change?
Type of change
Tested on
Testing instructions
Needs the
v0.10.0XNNPACK models re-exported from@bh/new-schemaon export-scripts, which swapsget_dynamic_dims_forwardfor aget_model_schemacompanion. The MLX models are unchanged.cd apps/nlp && yarn ios(oryarn android).Related issues
Closes #1246
Checklist
Additional notes
Viterbi decode runs in
O(numLabels + numEntities)per token instead ofO(numLabels²), cutting a ~27% CPU overhead on Nemotron's large (221) label space.