Add pplx-pii-masking

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Perplexity AI, Inc.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,73 @@
+---
+license: mit
+language:
+- en
+- multilingual
+tags:
+- pii
+- token-classification
+- privacy
+- qwen3
+pipeline_tag: token-classification
+---
+
+# pplx-pii-masking
+
+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, `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
+
+[`example_usage.py`](example_usage.py) in this repo is a self-contained
+reference pipeline (`pip install torch safetensors transformers`). It loads
+the encoder implementation from the
+[backbone repo](https://huggingface.co/perplexity-ai/pplx-embed-v1-0.6b) via
+`trust_remote_code`, swaps in this repo's fine-tuned weights, applies the two
+heads, and decodes spans with the constrained BIOES Viterbi included in the
+script:
+
+```python
+import sys
+
+from huggingface_hub import snapshot_download
+
+repo = snapshot_download("perplexity-ai/pplx-pii-masking")
+sys.path.insert(0, repo)
+from example_usage import PiiMasker
+
+masker = PiiMasker(repo)
+
+text = ("Hi, I'm Daniel Whitfield, you can reach me at "
+        "daniels@meridiancap.com or 415-555-0123.")
+spans, sensitivity = masker(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(masker.mask(text))
+# Hi, I'm [PRIVATE_PERSON], you can reach me at [PRIVATE_EMAIL] or [PRIVATE_PHONE].
+```
+
+## 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.py` / `modeling.py`
+referenced by `config.json`'s `auto_map`) ships with the
+[backbone repo](https://huggingface.co/perplexity-ai/pplx-embed-v1-0.6b).

config.json

@@ -0,0 +1,95 @@
+{
+  "model_type": "pii_masking",
+  "backbone": {
+    "vocab_size": 151936,
+    "max_position_embeddings": 32768,
+    "hidden_size": 1024,
+    "intermediate_size": 3072,
+    "num_hidden_layers": 28,
+    "num_attention_heads": 16,
+    "use_sliding_window": false,
+    "sliding_window": null,
+    "max_window_layers": 28,
+    "num_key_value_heads": 8,
+    "head_dim": 128,
+    "hidden_act": "silu",
+    "initializer_range": 0.02,
+    "rms_norm_eps": 1e-06,
+    "use_cache": false,
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "layer_types": [
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention"
+    ],
+    "pad_token_id": null,
+    "bos_token_id": 151643,
+    "eos_token_id": 151643,
+    "tie_word_embeddings": true,
+    "rope_parameters": {
+      "rope_theta": 1000000,
+      "rope_type": "default"
+    },
+    "return_dict": true,
+    "output_hidden_states": false,
+    "dtype": "bfloat16",
+    "chunk_size_feed_forward": 0,
+    "is_encoder_decoder": false,
+    "architectures": [
+      "PPLXQwen3Model"
+    ],
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "problem_type": null,
+    "_name_or_path": "perplexity-ai/pplx-embed-v1-0.6b",
+    "transformers_version": "5.2.0",
+    "auto_map": {
+      "AutoConfig": "configuration.PPLXQwen3Config",
+      "AutoModel": "modeling.PPLXQwen3Model"
+    },
+    "model_type": "bidirectional_pplx_qwen3",
+    "use_bidirectional_attention": true,
+    "output_attentions": false
+  },
+  "num_token_labels": 37,
+  "hidden_size": 1024,
+  "loss_weight_token_cls": 1.5,
+  "loss_weight_sensitivity": 0.3,
+  "dropout": 0.1,
+  "max_seq_len": 4096,
+  "viterbi_b_bias": 0.0,
+  "viterbi_e_bias": 0.0,
+  "freeze_backbone": false
+}

example_usage.py

@@ -0,0 +1,275 @@
+#!/usr/bin/env python3
+"""Detect and mask PII with pplx-pii-masking.
+
+Self-contained: the encoder implementation (`PPLXQwen3Model`) is pulled from
+the public backbone repo perplexity-ai/pplx-embed-v1-0.6b via
+`trust_remote_code`, the fine-tuned weights come from this repo, and the
+constrained BIOES Viterbi decoder is inlined below.
+
+  pip install torch safetensors transformers
+  python example_usage.py "some text with daniels@meridiancap.com in it"
+"""
+from __future__ import annotations
+
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file
+from transformers import AutoModel, AutoTokenizer
+
+BACKBONE_REPO = "perplexity-ai/pplx-embed-v1-0.6b"
+
+PII_TYPES = [
+    "private_person",
+    "private_email",
+    "private_phone",
+    "private_address",
+    "private_url",
+    "private_date",
+    "account_number",
+    "secret",
+    "other_pii",
+]
+# BIOES label list: O + 9 types x {B, I, E, S} = 37 (matches the checkpoint)
+BIOES_LABELS = ["O"] + [f"{tag}-{t}" for t in PII_TYPES for tag in "BIES"]
+
+
+# ---------------------------------------------------------------------------
+# Constrained BIOES Viterbi decoder
+# ---------------------------------------------------------------------------
+
+@dataclass
+class PredictedSpan:
+    start: int   # character offset in document text
+    end: int
+    label: str   # e.g. "private_person"
+    score: float
+
+
+def strip_span_whitespace(text: str, spans: list[PredictedSpan]) -> list[PredictedSpan]:
+    """Trim leading/trailing whitespace from predicted character spans.
+
+    Tokenizers fuse the space before a word into the token (e.g. '_john'),
+    so the B-token's char_start is one position before the actual PII value.
+    """
+    out: list[PredictedSpan] = []
+    for s in spans:
+        new_start, new_end = s.start, s.end
+        while new_start < new_end and text[new_start] in (" ", "\t", "\n"):
+            new_start += 1
+        while new_end > new_start and text[new_end - 1] in (" ", "\t", "\n"):
+            new_end -= 1
+        if new_start < new_end:
+            out.append(PredictedSpan(new_start, new_end, s.label, s.score))
+    return out
+
+
+class ViterbiDecoder:
+    """Constrained BIOES Viterbi decoder with transition bias scalars.
+
+    The bias scalars are added to all ->B and E-> transitions, allowing
+    precision/recall trade-offs without retraining.
+    """
+
+    def __init__(self, labels: list[str], b_bias: float = 0.0, e_bias: float = 0.0):
+        self.labels = labels
+        self.num_labels = len(labels)
+        self.label2id = {label: idx for idx, label in enumerate(labels)}
+        self.id2label = {idx: label for idx, label in enumerate(labels)}
+        self.pii_types = [label[2:] for label in labels if label.startswith("S-")]
+        self.b_bias = b_bias
+        self.e_bias = e_bias
+
+    def _build_transition_scores(self) -> torch.Tensor:
+        """[num_labels, num_labels] float transition score matrix."""
+        n, l2i = self.num_labels, self.label2id
+        mask = torch.zeros(n, n, dtype=torch.bool)
+        end_states = {l2i["O"]}    # states a span can end on (O, E-*, S-*)
+        begin_states = {l2i["O"]}  # states valid after a boundary (O, B-*, S-*)
+        for pii_type in self.pii_types:
+            b, i = l2i[f"B-{pii_type}"], l2i[f"I-{pii_type}"]
+            e, s = l2i[f"E-{pii_type}"], l2i[f"S-{pii_type}"]
+            end_states |= {e, s}
+            begin_states |= {b, s}
+            mask[b, i] = mask[b, e] = True   # B -> I/E of same type
+            mask[i, i] = mask[i, e] = True   # I -> I/E of same type
+        for from_state in end_states:
+            for to_state in begin_states:
+                mask[from_state, to_state] = True
+
+        scores = torch.full((n, n), float("-inf"))
+        scores[mask] = 0.0
+        for pii_type in self.pii_types:
+            b, e = l2i[f"B-{pii_type}"], l2i[f"E-{pii_type}"]
+            for from_s in range(n):
+                if scores[from_s, b] > float("-inf"):
+                    scores[from_s, b] += self.b_bias   # entering B states
+            for to_s in range(n):
+                if scores[e, to_s] > float("-inf"):
+                    scores[e, to_s] += self.e_bias     # leaving E states
+        return scores
+
+    @torch.no_grad()
+    def decode(
+        self,
+        logits: torch.Tensor,
+        offset_mapping: list[tuple[int, int]],
+        text: str | None = None,
+    ) -> list[PredictedSpan]:
+        """Decode a single sequence.
+
+        Args:
+            logits: [T, num_labels] float tensor
+            offset_mapping: list of (char_start, char_end) per token
+            text: source text; when provided, leading/trailing whitespace is
+                stripped from predicted span boundaries.
+        """
+        T, C = logits.shape
+        assert C == self.num_labels
+
+        trans = self._build_transition_scores()  # [C, C]
+
+        viterbi_scores = torch.full((T, C), float("-inf"))
+        backpointers = torch.zeros((T, C), dtype=torch.long)
+
+        # t=0: only O / B-* / S-* are valid start states
+        start_mask = torch.full((C,), float("-inf"))
+        for label, idx in self.label2id.items():
+            if label == "O" or label.startswith("B-") or label.startswith("S-"):
+                start_mask[idx] = 0.0
+        viterbi_scores[0] = logits[0] + start_mask
+
+        for t in range(1, T):
+            # [C, 1] + [C, C] -> [C, C]; dim-0 = prev, dim-1 = next
+            scores_t = viterbi_scores[t - 1].unsqueeze(1) + trans
+            best_prev, best_idx = scores_t.max(dim=0)
+            viterbi_scores[t] = logits[t] + best_prev
+            backpointers[t] = best_idx
+
+        # End constraint: only O / E-* / S-* valid at end
+        end_mask = torch.full((C,), float("-inf"))
+        for label, idx in self.label2id.items():
+            if label == "O" or label.startswith("E-") or label.startswith("S-"):
+                end_mask[idx] = 0.0
+        best_last = int((viterbi_scores[T - 1] + end_mask).argmax().item())
+
+        path = [best_last]
+        for t in range(T - 1, 0, -1):
+            path.append(int(backpointers[t, path[-1]].item()))
+        path.reverse()
+
+        spans = self._path_to_spans(path, logits, offset_mapping)
+        if text is not None:
+            spans = strip_span_whitespace(text, spans)
+        return spans
+
+    def _path_to_spans(
+        self,
+        path: list[int],
+        logits: torch.Tensor,
+        offset_mapping: list[tuple[int, int]],
+    ) -> list[PredictedSpan]:
+        spans: list[PredictedSpan] = []
+        T = len(path)
+        t = 0
+        while t < T:
+            label = self.id2label[path[t]]
+            char_start, char_end = offset_mapping[t]
+
+            if label.startswith("S-"):
+                pii_type = label[2:]
+                score = float(logits[t, path[t]].sigmoid().item())
+                if char_start < char_end:
+                    spans.append(PredictedSpan(char_start, char_end, pii_type, score))
+                t += 1
+
+            elif label.startswith("B-"):
+                pii_type = label[2:]
+                span_start, span_end = char_start, char_end
+                tok_scores = [float(logits[t, path[t]].item())]
+                t += 1
+                while t < T:
+                    inner = self.id2label[path[t]]
+                    if inner == f"I-{pii_type}":
+                        _, span_end = offset_mapping[t]
+                        tok_scores.append(float(logits[t, path[t]].item()))
+                        t += 1
+                    elif inner == f"E-{pii_type}":
+                        _, span_end = offset_mapping[t]
+                        tok_scores.append(float(logits[t, path[t]].item()))
+                        t += 1
+                        break
+                    else:
+                        break
+                if span_start < span_end:
+                    score = float(torch.tensor(tok_scores).mean().sigmoid().item())
+                    spans.append(PredictedSpan(span_start, span_end, pii_type, score))
+
+            else:
+                t += 1
+
+        return spans
+
+
+# ---------------------------------------------------------------------------
+# Model loading and inference
+# ---------------------------------------------------------------------------
+
+class PiiMasker:
+    def __init__(self, repo_dir: str | Path, device: str = "cpu"):
+        repo_dir = Path(repo_dir)
+        self.device = device
+        self.tokenizer = AutoTokenizer.from_pretrained(repo_dir)
+
+        # Encoder architecture from the public backbone repo; weights are then
+        # replaced with the fine-tuned backbone from this checkpoint.
+        self.backbone = AutoModel.from_pretrained(
+            BACKBONE_REPO, dtype=torch.bfloat16, trust_remote_code=True
+        ).to(device).eval()
+
+        sd = load_file(str(repo_dir / "model.safetensors"))
+        self.backbone.load_state_dict(
+            {k.removeprefix("backbone."): v for k, v in sd.items()
+             if k.startswith("backbone.")}
+        )
+        self.w_cls = sd["token_cls_head.weight"].float().to(device)   # [37, 1024]
+        self.b_cls = sd["token_cls_head.bias"].float().to(device)
+        self.w_sen = sd["sensitivity_head.weight"].float().to(device)  # [1, 1024]
+        self.b_sen = sd["sensitivity_head.bias"].float().to(device)
+        self.viterbi = ViterbiDecoder(
+            BIOES_LABELS,
+            b_bias=float(sd["viterbi.b_bias"].item()),
+            e_bias=float(sd["viterbi.e_bias"].item()),
+        )
+
+    @torch.no_grad()
+    def __call__(self, text: str) -> tuple[list[PredictedSpan], float]:
+        enc = self.tokenizer(text, return_offsets_mapping=True,
+                             return_tensors="pt", truncation=True, max_length=4096)
+        h = self.backbone(
+            input_ids=enc["input_ids"].to(self.device),
+            attention_mask=enc["attention_mask"].to(self.device),
+        ).last_hidden_state[0].float()                     # [T, 1024]
+        logits = h @ self.w_cls.T + self.b_cls             # [T, 37]
+        sensitivity = torch.sigmoid(h.mean(0) @ self.w_sen.T + self.b_sen).item()
+        offsets = [tuple(o) for o in enc["offset_mapping"][0].tolist()]
+        spans = self.viterbi.decode(logits.cpu(), offsets, text=text)
+        return spans, sensitivity
+
+    def mask(self, text: str) -> str:
+        spans, _ = self(text)
+        for s in sorted(spans, key=lambda s: -s.start):
+            text = text[:s.start] + f"[{s.label.upper()}]" + text[s.end:]
+        return text
+
+
+if __name__ == "__main__":
+    text = sys.argv[1]
+    masker = PiiMasker(Path(__file__).parent)
+    spans, sensitivity = masker(text)
+    print(f"sensitivity: {sensitivity:.3f}")
+    for s in spans:
+        print(f"  {s.label:18s} [{s.start}:{s.end}] {text[s.start:s.end]!r}")
+    print(masker.mask(text))

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f6204155ec540c9323f706e284110ee848b462f0325dc1ece5c7263fc517bbd0
+size 1192293777

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cae14d1c8dda080f23792355b0692b826bf1f1da3c86ebc1b37548a391cf6526
+size 11422936

tokenizer_config.json

@@ -0,0 +1,15 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "mask_token": "â½Ĺ",
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "sep_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "unk_token": null
+}