pplx-pii-masking

PII masking model for conversational data. A ~600M-parameter bidirectional Qwen3 encoder (perplexity-ai/pplx-embed-v1-0.6b backbone, use_bidirectional_attention=true) 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, then load the model with trust_remote_code:

from transformers import AutoModel

model = AutoModel.from_pretrained(
    "perplexity-ai/pplx-pii-masking", 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 from modeling_pplx_qwen3.py (vendored from the backbone repo), 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 PPLXQwen3Model encoder implementation (configuration_pplx_qwen3.py / modeling_pplx_qwen3.py) is vendored from the backbone repo; config.json's backbone holds its config.