PII-Tracer-MLX
This repository provides PII-Tracer, the detector introduced in PII-TRACE.
MLX weights for PII-Tracer,
a PII masking model for conversational data. The model uses a ~600M-parameter
bidirectional Qwen3 encoder with a
perplexity-ai/pplx-embed-v1-0.6b
backbone and two heads:
- Token classification head (1024 -> 37): BIOES tags over 9 PII categories
(
private_person,account_number,private_url,private_date,private_address,private_email,private_phone,other_pii,secret), decoded with a constrained Viterbi decoder. - Sensitivity head (1024 -> 1): conversation-level sensitivity classifier on mean-pooled hidden states.
The backbone's linear projections use 8-bit affine quantization with group size 64. Embeddings and normalization weights remain bf16; both heads remain fp32. The weights file is approximately 779 MB.
Contents
| File | Description |
|---|---|
weights.safetensors | Quantized backbone and both heads |
config.json | Backbone configuration, quantization parameters, 4096-token limit, and Viterbi biases |
tokenizer.json, tokenizer_config.json | Tokenizer and character-offset mapping |
LICENSE | MIT license, matching the base model |
Quantized projections retain the backbone.* key names and store packed
uint32 weight tensors with bf16 scales and biases. The unquantized heads
use token_cls_head.* and sensitivity_head.*.
Inference integration
This repository supplies converted weights and configuration. It does not
include an MLX model implementation or loading script. Using these weights
requires an MLX implementation of the bidirectional Qwen3 encoder, affine
quantized projections, and the two heads. It must honor
config.json's backbone.is_causal = false and mask padding during attention
and mean pooling.
Tokenize without adding BOS/EOS tokens. The model's input limit is 4096
tokens; chunk longer documents before inference. For encoder hidden states
hidden_states, compute:
token_logits = hidden_states @ token_cls_head.weight.T + token_cls_head.bias
sensitivity_logit = mean_pool(hidden_states) @ sensitivity_head.weight.T + sensitivity_head.bias
sensitivity = sigmoid(sensitivity_logit)
Here mean_pool excludes padding tokens. Decode token_logits with a
constrained BIOES Viterbi decoder using the biases in config.json, then map
token spans to character offsets. Label index 0 is O; each category has
four consecutive labels in B, I, E, S order, with categories ordered as:
private_person, private_email, private_phone, private_address, private_url,
private_date, account_number, secret, other_pii
The base model provides the reference model behavior and constrained BIOES decoding logic.
Citation
If you use or reference this work, please cite:
@article{zhang2026piitrace,
title = {{PII-TRACE}: A Benchmark for Context-Aware {PII} Detection in Multi-Turn {LLM} Conversations},
author = {Zhang, Kaiyuan and Wang, Chuan and Zhong, Joey and Fryzel, Paul and Polley, Kyle and Ma, Jerry and Li, Ninghui},
journal = {arXiv preprint arXiv:2609.22200},
year = {2026},
url = {https://arxiv.org/abs/2609.22200}
}