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:
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 [blocked] in this repo: it builds
the encoder from modeling_pplx_qwen3.py [blocked] (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.