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}
}