---
license: mit
language:
- en
- multilingual
library_name: mlx
tags:
- pii
- token-classification
- privacy
- qwen3
- mlx
- apple-silicon
pipeline_tag: token-classification
base_model: perplexity-ai/PII-Tracer
---

# PII-Tracer-MLX

This repository provides PII-Tracer, the detector introduced in
[PII-TRACE](https://arxiv.org/abs/2609.22200).

MLX weights for [PII-Tracer](https://huggingface.co/perplexity-ai/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`](https://huggingface.co/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:

```text
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:

```text
private_person, private_email, private_phone, private_address, private_url,
private_date, account_number, secret, other_pii
```

The [base model](https://huggingface.co/perplexity-ai/PII-Tracer)
provides the reference model behavior and constrained BIOES decoding logic.

## Citation

If you use or reference this work, please cite:

```bibtex
@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}
}
```