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