perplexity-ai

/

pplx-pii-masking

like 41

Follow

Perplexity 2.78k

Hello @kaiyuanzh and team!

Thanks for open-weighting this! I wanted to make it a tad simpler to load and use the model, so I worked on a simple custom code implementation for AutoModel. Heads up, this PR was AI-generated and human-reviewed.

Pull Request overview

• Add an AutoModel custom code implementation, so the model loads with trust_remote_code=True
• Rename example_usage.py to modeling_pii_masking.py and point config.json's auto_map at it
• Update the README usage snippet

Details

Using this model today means calling snapshot_download, inserting the result into sys.path and importing PiiMasker from example_usage.py. Everything needed for the normal trust_remote_code route was already in the repo, so this moves the reference pipeline behind it:

from transformers import AutoModel

model = AutoModel.from_pretrained("perplexity-ai/pplx-pii-masking", trust_remote_code=True)

spans, sensitivity = model.predict(text)
print(model.mask(text))

Most of the diff is the rename. PredictedSpan, strip_span_whitespace and the Viterbi decoding are unchanged, PiiMasker becomes a PiiMaskingConfig plus a PiiMaskingModel, and predict and mask keep the exact signatures and return values that PiiMasker.__call__ and PiiMasker.mask had. config.json picks up architectures, auto_map and dtype.

Three parts are worth a closer look:

ViterbiDecoder is now an nn.Module with b_bias, e_bias and transition_mask as buffers, so all three load from model.safetensors rather than being rebuilt in Python. viterbi.transition_mask was not read at all before. It is exactly equal to the mask the decoder derives from the label list (I checked), so nothing changes numerically. It does mean the checkpoint stays the source of truth, and that there are no unexpected keys at load time.

The encoder is built from config.backbone with get_class_from_dynamic_module instead of AutoModel.from_pretrained(BACKBONE_REPO). The old path downloaded the backbone's 1.2 GB model.safetensors and then replaced every tensor in it via load_state_dict. On a cold cache the first load now fetches only configuration.py and modeling.py from the backbone repo, about 7 KiB.

_keep_in_fp32_modules_strict keeps token_cls_head, sensitivity_head and the two decoder biases in fp32 while the encoder loads in bf16, matching how the checkpoint stores them and what the script did by hand with .float().

I compared this against the current example_usage.py on the README example, a multi category conversation, text with no PII, inputs above and below max_seq_len, and whitespace only and empty documents. Hidden states are bitwise identical, token logits differ by at most 1e-6 (nn.Linear uses addmm where the script wrote h @ W_cls.T + b_cls), and sensitivity agrees to nine decimals. Spans, span scores and masked text are identical in every case, including with non-zero viterbi_b_bias and viterbi_e_bias. save_pretrained round trips config, weights and the .py file, and loading with dtype=torch.float32 gives the same spans.

One deliberate behaviour change: model.predict("") returns ([], 0.0) where the old code raised, because an empty document tokenizes to zero tokens and the attention mask builder cannot take a zero length sequence.

Three things I would happily change:

I removed example_usage.py instead of keeping it as a wrapper. If you would rather keep the file and its CLI, it is about ten lines on top of model.predict.

The backbone code is fetched unpinned, as before. get_class_from_dynamic_module takes a code_revision, so the backbone commit could be pinned for reproducible loads.

I mapped AutoModel only. Adding AutoModelForTokenClassification would let pipeline("token-classification") load the model, but the pipeline decodes with argmax rather than the constrained Viterbi, so its output would not match predict.

• Tom Aarsen