PII-Tracer-vLLM
This repository provides PII-Tracer, the detector introduced in PII-TRACE.
vLLM packaging of
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
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.
This repository repackages the fine-tuned weights for vLLM without retraining
or quantization. It also includes a Docker Compose deployment, a /v1/scoring
adapter, a decoding client, and a reference forward implementation.
Usage
The Docker Compose deployment requires an NVIDIA GPU, a compatible driver, and Docker Compose with GPU support. It pins vLLM 0.26.0. Run these commands from the same shell:
pip install huggingface_hub numpy tokenizers requests
hf download perplexity-ai/PII-Tracer-vLLM --local-dir PII-Tracer-vLLM
cd PII-Tracer-vLLM
docker compose -f serving/docker-compose.yml up -d
curl -s http://127.0.0.1:8002/health
curl -s http://127.0.0.1:8002/v1/scoring -H 'Content-Type: application/json' \
-d '{"model":"pii-masking-latest","sequences":["My name is John Smith."]}'
python serving/test_client.py --model-dir . \
"Hi, I'm Daniel Whitfield, you can reach me at daniels@meridiancap.com or 415-555-0123."
Wait for both services to become healthy before sending requests. The client prints PII spans, masked text, and the sigmoid of the sensitivity logit. The example client expects nonempty text.
To run vLLM directly instead of using Docker Compose:
vllm serve perplexity-ai/PII-Tracer-vLLM \
--served-model-name pii-masker --host 127.0.0.1 --port 8003 \
--runner pooling --convert classify \
--pooler-config '{"use_activation": false}' \
--dtype bfloat16 --max-model-len 4096
This direct command starts the /pooling endpoint. The /v1/scoring endpoint
and decoding client use the additional adapter provided by Docker Compose.
POST /pooling with {"model": "pii-masker", "task": "token_classify", "input": [[token ids]]}
returns a [T, 38] logit matrix for each input. Tokenize without adding BOS/EOS.
Checkpoint layout
| Base checkpoint | vLLM packaging |
|---|---|
| Bidirectional Qwen3 encoder and two separate heads | Qwen3ForTokenClassification architecture with is_causal: false |
backbone.* | model.* |
token_cls_head.* (37 outputs) | Rows 0–36 of score.weight and score.bias |
sensitivity_head.* (1 output) | Row 37 of score.weight and score.bias |
viterbi.* | viterbi.json; the decoding client constructs the matching constraints |
The backbone weights retain bf16 precision and the head weights retain fp32
precision. The adapter separates the 37 token label logits from column 37,
then averages column 37 over the input tokens to obtain the sequence
sensitivity logit. This follows from
mean(W·h + b) = W·mean(h) + b; apply sigmoid to obtain sensitivity.
Floating-point results can vary with the serving dtype and backend.
Labels and decoding
The 37 token labels use index 0 = O. For category index c from 0 to 8,
1+4c = B, 2+4c = I, 3+4c = E, and 4+4c = S.
The checkpoint's category index order is:
private_person, private_email, private_phone, private_address,
private_url, private_date, account_number, secret, other_pii.
Column 37 is __SEQ_SENSITIVITY__ and is excluded from token decoding.
The client uses constrained BIOES Viterbi decoding and token offsets to
recover character spans. Its span confidence is the mean softmax probability
of the selected token labels; this differs from the base model's span score.
Limitations
- The input limit is 4096 tokens. The supplied tokenizer truncates longer input, so chunk longer documents before scoring to cover the entire text.
use_activation: falseis required for raw logits, Viterbi decoding, and recovery of the sequence sensitivity logit.- Span boundaries follow token offsets and can include surrounding whitespace. False positives and missed PII are possible; evaluate on your target data.
- Calibrate sensitivity and span confidence thresholds on your target data before using them to make masking decisions.
- Docker Compose binds both services to
127.0.0.1; the adapter does not enforce authentication. Adjust--gpu-memory-utilizationin the Compose file if the default allocation is insufficient for your GPU.
Files
| Path | Purpose |
|---|---|
config.json, model.safetensors, tokenizer.json, tokenizer_config.json, viterbi.json | Model, tokenizer, and decoding configuration |
serving/docker-compose.yml | vLLM and scoring adapter deployment |
serving/scoring_adapter.py | /v1/scoring adapter for /pooling |
serving/test_client.py | Viterbi decoding and masking example |
serving/convert_to_hf.py | Converter from the base checkpoint |
serving/reference_forward.py | fp32 CPU reference for comparison with a running scoring adapter |
License
MIT, matching the base model. See LICENSE.
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}
}