PII-Tracer
This repository provides PII-Tracer, the detector introduced in PII-TRACE.
PII masking model for conversational data. A ~600M-parameter bidirectional
Qwen3 encoder
(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:
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 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 repackages the weights for token classification serving, with a scoring adapter and a decoding client.
- GGUF: 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 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:
@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}
}