PII-Tracer

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