---
license: mit
language:
- en
- multilingual
tags:
- pii
- token-classification
- privacy
- qwen3
pipeline_tag: token-classification
---

# PII-Tracer

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

PII masking model for conversational data. A ~600M-parameter bidirectional
Qwen3 encoder
([`perplexity-ai/pplx-embed-v1-0.6b`](https://huggingface.co/perplexity-ai/pplx-embed-v1-0.6b)
backbone, run encoder-only via `is_causal=false`) with 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.

## Usage

`pip install torch "transformers>=5.2"`, then load the model with
`trust_remote_code`:

```python
from transformers import AutoModel

model = AutoModel.from_pretrained(
    "perplexity-ai/PII-Tracer", trust_remote_code=True
)

text = ("Hi, I'm Daniel Whitfield, you can reach me at "
        "daniels@meridiancap.com or 415-555-0123.")
spans, sensitivity = model.predict(text)
for s in spans:
    print(s.label, (s.start, s.end), text[s.start:s.end])
# private_person (8, 24) Daniel Whitfield
# private_email (46, 69) daniels@meridiancap.com
# private_phone (73, 85) 415-555-0123

print(model.mask(text))
# Hi, I'm [PRIVATE_PERSON], you can reach me at [PRIVATE_EMAIL] or [PRIVATE_PHONE].
```

`predict` and `mask` wrap `model(input_ids, attention_mask)`, which returns the
37 BIOES tag logits per token and one sensitivity logit per document. The
implementation is
[`modeling_pii_masking.py`](modeling_pii_masking.py) in this repo: it builds
the encoder with the standard Transformers `Qwen3Model` run bidirectionally
via `config.is_causal = False` (hence `transformers>=5.2`, which passes
`is_causal` through to the attention backend),
applies this repo's fine-tuned weights and the two heads, and decodes spans
with the constrained BIOES Viterbi included in the file. Input is truncated
to `max_seq_len` (4096) tokens; chunk longer documents before calling
`predict`. `model.save_pretrained(dir)` also saves the tokenizer, so the saved
directory loads and predicts on its own.

## Checkpoint layout

`model.safetensors` holds the fine-tuned backbone (bf16, `backbone.*`), both
heads (fp32, `token_cls_head.*` / `sensitivity_head.*`), and the Viterbi bias
scalars (`viterbi.*`). `max_seq_len` is 4096 tokens.

Inference outline: tokenize (no BOS/EOS added), run the bidirectional encoder,
then per token `logits = h @ W_cls.T + b_cls` decoded with a constrained BIOES
Viterbi, and `sensitivity = sigmoid(mean(h) @ W_sen.T + b_sen)`. The encoder
is the stock Transformers `Qwen3Model`; `config.json`'s `backbone` holds its
config, including `"is_causal": false` for bidirectional attention.

## Alternative packagings

- **vLLM:** [perplexity-ai/PII-Tracer-vLLM](https://huggingface.co/perplexity-ai/PII-Tracer-vLLM)
  repackages the weights for token classification serving, with a scoring
  adapter and a decoding client.
- **GGUF:** [perplexity-ai/PII-Tracer-GGUF](https://huggingface.co/perplexity-ai/PII-Tracer-GGUF)
  provides an f16 backbone for llama-server, fp32 heads, and a self-contained
  decoding client.
- **MLX:** [perplexity-ai/PII-Tracer-MLX](https://huggingface.co/perplexity-ai/PII-Tracer-MLX)
  provides 8-bit converted projection weights for MLX. An MLX implementation
  of the bidirectional encoder and two heads is required; a loader is not
  included in that repository.

See each repository for its inference requirements and limitations.

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