lemonade-sdk/pplx-pii-masking-onnx

pplx-pii-masking - ONNX

This is an ONNX export of `perplexity-ai/pplx-pii-masking`, converted for CPU inference with `onnxruntime`'s CPUExecutionProvider. No fine-tuning was performed, this repo only changes the runtime format, not the model's weights or behavior.

About the base model

All credit for the model itself goes to [Perplexity](https://huggingface.co/perplexity-ai). From their model card: pplx-pii-masking is a ~600M-parameter bidirectional Qwen3 encoder (the `pplx-embed-v1-0.6b` backbone with use_bidirectional_attention=true) fine-tuned for PII masking in conversational data, with two heads on the shared backbone:

• —Token classification head (1024 -> 37): BIOES tags over 9 PII categories - private_person, private_email, private_phone, private_address, private_url, private_date, account_number, secret, other_pii - decoded upstream with a constrained Viterbi decoder.
• —Sensitivity head (1024 -> 1): a document-level sensitivity score over mean-pooled hidden states.

See the original model card for training details and the full inference outline.

What was done for this conversion

Exported from the original repo. This export loads the original repo with trust_remote_code=True, whose vendored modeling_pplx_qwen3.py gets it right on purpose: post_init() flips self_attn.is_causal = False on every layer and forward() rebuilds a real padding-only bidirectional mask.

1. 1.Both heads are exported, as two named outputs. logits is [batch, seq_len, 37] (per-token BIOES logits) and sensitivity_logits is [batch] (one pre-sigmoid scalar per sequence).
2. 2.*The constrained BIOES Viterbi decoder is not baked into the graph. Upstream keeps it outside `forward()` (it lives in `modeling_pii_masking.py`'s `ViterbiDecoder` and runs inside `model.predict()`), and this export keeps the same boundary: the graph stops at raw per-token logits so downstream consumers can apply the checkpoint's own decoder, their own decoder, or a threshold scheme. See Decoding* below for what that means in practice.
3. 3.fp32 throughout. config.json declares a bf16 default and AutoModel honors it unless overridden. A bf16 export produces Where nodes with bf16 scalar operands that onnxruntime's CPU EP has no kernel for - a NOT_IMPLEMENTED error at session-load time, not export time. This export loads and traces in full fp32.

config.json, tokenizer.json and tokenizer_config.json are copied verbatim from the original repo. Note that config.json's auto_map entries reference modeling_pii_masking.py, which is not shipped here - the tokenizer loads without it, and onnxruntime consumers only need architectures / max_seq_len from that file.

Verification

The exported graph was validated against an eager fp32 PyTorch forward pass of the original checkpoint, on onnxruntime's CPUExecutionProvider:

• —Max absolute logit difference (ONNX vs. PyTorch), canned sample: 1.5e-5 (logits), 2e-6 (sensitivity_logits)
• —Max absolute logit difference, real 806-character document: 1.0e-5 (logits), 0.0 (sensitivity_logits)
• —Decision parity on 20,000 PII-bearing documents (nvidia/Nemotron-PII), scoring both backends through the checkpoint's own ViterbiDecoder: identical confusion matrix, 100.0000% has-PII agreement, identical per-document label sets, and the same 159 missed documents by name - set difference empty both ways.

Usage

import numpy as np
import onnxruntime as ort
from transformers import AutoTokenizer

REPO = "lemonade-sdk/pplx-pii-masking-onnx"
tokenizer = AutoTokenizer.from_pretrained(REPO)
session = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])

text = "Hi, I'm Daniel Whitfield, you can reach me at daniel@meridiancap.com or 415-555-0123."
enc = tokenizer(text, return_tensors="np", truncation=True, max_length=4096)
logits, sensitivity_logits = session.run(
    ["logits", "sensitivity_logits"],
    {
        "input_ids": enc["input_ids"].astype(np.int64),
        "attention_mask": enc["attention_mask"].astype(np.int64),
    },
)
# logits: [1, seq_len, 37]   sensitivity_logits: [1]
sensitivity = 1 / (1 + np.exp(-sensitivity_logits[0]))

The graph takes a plain 2D input_ids + attention_mask, both int64 - recent transformers versions hand back int32 from return_tensors="np", hence the cast. No BOS/EOS is added by the tokenizer, matching upstream. Input is truncated at max_seq_len (4096); chunk longer documents.

Decoding

The 37 labels are ["O"] + [f"{tag}-{t}" for t in PII_TYPES for tag in "BIES"], in that order - this is the id2label in manifest.json. To reproduce model.predict() exactly, apply the checkpoint's own decoder to the ONNX logits. Its whole state is the label list plus two bias scalars from config.json, so no weights need loading:

import importlib
import torch
from transformers import AutoConfig
from transformers.dynamic_module_utils import get_class_from_dynamic_module

cfg = AutoConfig.from_pretrained("perplexity-ai/pplx-pii-masking", trust_remote_code=True)
ViterbiDecoder = get_class_from_dynamic_module(
    "modeling_pii_masking.ViterbiDecoder", "perplexity-ai/pplx-pii-masking")
modeling = importlib.import_module(ViterbiDecoder.__module__)
decoder = ViterbiDecoder(modeling.BIOES_LABELS, cfg.viterbi_b_bias, cfg.viterbi_e_bias).eval()

enc = tokenizer(text, return_offsets_mapping=True, truncation=True, max_length=4096)
spans = decoder.decode(torch.from_numpy(logits[0]).float(), enc["offset_mapping"], text=text)
for s in spans:
    print(s.label, text[s.start:s.end])

On the two heads: they disagree in the normal case, and the token head is the one to trust for detection. Measured over the 20,000-document benchmark above, the span head reached 99.20% recall; the sensitivity head at a 0.5 threshold reached 9.16%. The original model card's own example returns three correct spans at sensitivity=0.027. Do not gate PII detection on sensitivity_logits.

Using this with the Lemonade router

This repo is self-contained for Lemonade's onnxruntime backend (via `ort-server`) - model.onnx + model.onnx.data + tokenizer.json + config.json + manifest.json all sit together at the repo root.

Requires an `ort-server` build newer than 0.3.7.

1. Register the classifier:

curl -X POST http://localhost:13305/v1/pull -H "Content-Type: application/json" -d '{
  "model_name": "user.pplx-pii-masking-onnx",
  "checkpoint": "lemonade-sdk/pplx-pii-masking-onnx",
  "recipe": "onnxruntime"
}'

manifest.json declares task: token-classification, the 37-entry id2label, score_normalization: softmax and token_aggregation: max. It keeps all 37 raw BIOES labels distinct rather than collapsing B/I/E/S into one entity name, because ort-server writes per-label scores by plain assignment keyed on the label string - two indices sharing a name would silently overwrite each other rather than combine.

2. Register the router policy - a ready-to-use collection.router policy that routes to a local model whenever any of the 36 non-"O" labels crosses min_score: 0.5, and to a cloud model otherwise, is published separately at `lemonade-sdk/pii_policy_pplx-pii-masking-onnx`:

hf download lemonade-sdk/pii_policy_pplx-pii-masking-onnx --local-dir .
curl -X POST http://localhost:13305/v1/pull -H "Content-Type: application/json" \
  --data-binary @pii_policy_pplx-pii-masking-onnx.json

The policy's routing.candidates (Qwen3.5-0.8B-GGUF local / fireworks.kimi-k2p6 cloud) and min_score are starting points, not fixed requirements.

Note on the decision rule. The router does not run the Viterbi decoder - ort-server has no hook for it. Its rule is per-token softmax, then max over tokens per label, then "any non-O label >= min_score". That is a different decision rule from model.predict(), and it is not strictly looser or stricter: it can fire on a single confident token that the decoder's transition constraints would have vetoed, and it can miss a span the decoder would have assembled from several individually sub-threshold tokens. On the Nemotron-PII benchmark the two rules are close, but treat the router's leak rate as a property of the rule + threshold, not of the model alone.

For building your own routing policy from scratch, see the `lemonade-router-builder` skill.

License

MIT, inherited from the base model.

Citation

Please cite the original model:

@misc{pplx-pii-masking,
  title  = {pplx-pii-masking},
  author = {Perplexity AI},
  year   = {2026},
  url    = {https://huggingface.co/perplexity-ai/pplx-pii-masking}
}