Add an AutoModel custom code implementation

README.md

@@ -27,38 +27,41 @@ backbone, `use_bidirectional_attention=true`) with two heads:
 
 ## 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:
+`pip install torch transformers`, then load the model with `trust_remote_code`:
 
 ```python
-import sys
+from transformers import AutoModel
 
-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)
+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 = masker(text)
+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(masker.mask(text))
+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`](modeling_pii_masking.py) in this repo: it builds
+the encoder from [`modeling_pplx_qwen3.py`](modeling_pplx_qwen3.py) (vendored
+from the
+[backbone repo](https://huggingface.co/perplexity-ai/pplx-embed-v1-0.6b)),
+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
@@ -68,6 +71,7 @@ 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).
+`PPLXQwen3Model` encoder implementation (`configuration_pplx_qwen3.py` /
+`modeling_pplx_qwen3.py`) is vendored from the
+[backbone repo](https://huggingface.co/perplexity-ai/pplx-embed-v1-0.6b);
+`config.json`'s `backbone` holds its config.

config.json

@@ -1,4 +1,13 @@
 {
+  "architectures": [
+    "PiiMaskingModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "modeling_pii_masking.PiiMaskingConfig",
+    "AutoModel": "modeling_pii_masking.PiiMaskingModel"
+  },
+  "dtype": "bfloat16",
+  "transformers_version": "5.2.0",
   "model_type": "pii_masking",
   "backbone": {
     "vocab_size": 151936,
@@ -75,10 +84,6 @@
     "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
@@ -92,4 +97,4 @@
   "viterbi_b_bias": 0.0,
   "viterbi_e_bias": 0.0,
   "freeze_backbone": false
-}
+}

configuration_pplx_qwen3.py

@@ -0,0 +1,7 @@
+# Vendored from perplexity-ai/pplx-embed-v1-0.6b @ 2c4d510dd4a732063c31a0f70193e35067b51fd8
+# (configuration.py). Kept in this repo so loading needs no code from another repo.
+from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
+
+
+class PPLXQwen3Config(Qwen3Config):
+    model_type = "bidirectional_pplx_qwen3"

example_usage.py → modeling_pii_masking.py

@@ -1,25 +1,35 @@
-#!/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.
+Loaded through `AutoModel`, which pulls this file in as remote code:
 
-  pip install torch safetensors transformers
-  python example_usage.py "some text with daniels@meridiancap.com in it"
+    from transformers import AutoModel
+
+    model = AutoModel.from_pretrained(
+        "perplexity-ai/pplx-pii-masking", trust_remote_code=True
+    )
+    spans, sensitivity = model.predict("Hi, I'm Daniel Whitfield")
+    print(model.mask("Hi, I'm Daniel Whitfield"))
+
+The encoder implementation (`PPLXQwen3Model`) is vendored from the public
+backbone repo perplexity-ai/pplx-embed-v1-0.6b (`modeling_pplx_qwen3.py`), the
+fine-tuned weights come from this repo, and the constrained BIOES Viterbi
+decoder is inlined below.
 """
 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
+from torch import nn
+from transformers import AutoTokenizer, PretrainedConfig, PreTrainedModel
+from transformers.utils import ModelOutput
+
+from .configuration_pplx_qwen3 import PPLXQwen3Config
+from .modeling_pplx_qwen3 import PPLXQwen3Model
 
-BACKBONE_REPO = "perplexity-ai/pplx-embed-v1-0.6b"
+# from_pretrained kwargs that also apply when loading the tokenizer next to the weights
+HUB_KWARGS = ("cache_dir", "force_download", "local_files_only", "proxies",
+              "revision", "subfolder", "token")
 
 PII_TYPES = [
     "private_person",
@@ -66,24 +76,27 @@ def strip_span_whitespace(text: str, spans: list[PredictedSpan]) -> list[Predict
     return out
 
 
-class ViterbiDecoder:
+class ViterbiDecoder(nn.Module):
     """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.
+    precision/recall trade-offs without retraining. They and the legal
+    transition mask are buffers, so they come from the checkpoint.
     """
 
     def __init__(self, labels: list[str], b_bias: float = 0.0, e_bias: float = 0.0):
+        super().__init__()
         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
+        self.register_buffer("b_bias", torch.tensor([float(b_bias)]))
+        self.register_buffer("e_bias", torch.tensor([float(e_bias)]))
+        self.register_buffer("transition_mask", self._build_transition_mask())
 
-    def _build_transition_scores(self) -> torch.Tensor:
-        """[num_labels, num_labels] float transition score matrix."""
+    def _build_transition_mask(self) -> torch.Tensor:
+        """[num_labels, num_labels] bool mask of the legal tag transitions."""
         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-*)
@@ -98,17 +111,22 @@ class ViterbiDecoder:
         for from_state in end_states:
             for to_state in begin_states:
                 mask[from_state, to_state] = True
+        return mask
 
+    def _build_transition_scores(self) -> torch.Tensor:
+        """[num_labels, num_labels] float transition score matrix."""
+        n, l2i = self.num_labels, self.label2id
+        b_bias, e_bias = float(self.b_bias), float(self.e_bias)
         scores = torch.full((n, n), float("-inf"))
-        scores[mask] = 0.0
+        scores[self.transition_mask.bool().cpu()] = 0.0  # decoding runs on cpu
         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
+                    scores[from_s, b] += 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
+                    scores[e, to_s] += e_bias     # leaving E states
         return scores
 
     @torch.no_grad()
@@ -214,62 +232,156 @@ class ViterbiDecoder:
 
 
 # ---------------------------------------------------------------------------
-# Model loading and inference
+# Model
 # ---------------------------------------------------------------------------
 
-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.")}
+class PiiMaskingConfig(PretrainedConfig):
+    """Config for `PiiMaskingModel`.
+
+    `backbone` stays a plain dict; `PiiMaskingModel` turns it into a
+    `PPLXQwen3Config`.
+    """
+
+    model_type = "pii_masking"
+
+    def __init__(
+        self,
+        backbone: dict | None = None,
+        hidden_size: int = 1024,
+        num_token_labels: int = len(BIOES_LABELS),
+        max_seq_len: int = 4096,
+        viterbi_b_bias: float = 0.0,
+        viterbi_e_bias: float = 0.0,
+        **kwargs,
+    ):
+        self.backbone = backbone or {}
+        self.hidden_size = hidden_size
+        self.num_token_labels = num_token_labels
+        self.max_seq_len = max_seq_len
+        self.viterbi_b_bias = viterbi_b_bias
+        self.viterbi_e_bias = viterbi_e_bias
+        kwargs.setdefault("id2label", dict(enumerate(BIOES_LABELS)))
+        kwargs.setdefault(
+            "label2id", {label: i for i, label in enumerate(BIOES_LABELS)}
         )
-        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)
+        super().__init__(**kwargs)
+
+
+@dataclass
+class PiiMaskingOutput(ModelOutput):
+    """[B, T, 37] tag logits, [B] sensitivity logits, [B, T, H] encoder output."""
+
+    logits: torch.FloatTensor | None = None
+    sensitivity_logits: torch.FloatTensor | None = None
+    last_hidden_state: torch.FloatTensor | None = None
+
+
+class PiiMaskingModel(PreTrainedModel):
+    """Bidirectional Qwen3 encoder with PII token and sensitivity heads."""
+
+    config_class = PiiMaskingConfig
+    _supports_sdpa = True
+    _supports_flash_attn = True
+    # The heads and the decoder biases are fp32 in the checkpoint while the
+    # encoder is bf16. Keep them fp32 whichever dtype the model is loaded in.
+    _keep_in_fp32_modules_strict = [
+        "token_cls_head", "sensitivity_head", "viterbi.b_bias", "viterbi.e_bias",
+    ]
+
+    def __init__(self, config: PiiMaskingConfig):
+        super().__init__(config)
+        backbone_config = PPLXQwen3Config.from_dict(config.backbone)
+        backbone_config._attn_implementation = config._attn_implementation
+        self.backbone = PPLXQwen3Model(backbone_config)
+        self.token_cls_head = nn.Linear(config.hidden_size, config.num_token_labels)
+        self.sensitivity_head = nn.Linear(config.hidden_size, 1)
         self.viterbi = ViterbiDecoder(
-            BIOES_LABELS,
-            b_bias=float(sd["viterbi.b_bias"].item()),
-            e_bias=float(sd["viterbi.e_bias"].item()),
+            BIOES_LABELS, config.viterbi_b_bias, config.viterbi_e_bias
+        )
+        self._tokenizer = None
+        self._tokenizer_kwargs: dict = {}
+        self.post_init()
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        model = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+        # Same revision / token / cache as the weights, so `predict` on a model
+        # loaded from a branch or a private repo finds the matching tokenizer.
+        model._tokenizer_kwargs = {k: kwargs[k] for k in HUB_KWARGS if k in kwargs}
+        return model
+
+    def save_pretrained(self, save_directory, *args, **kwargs):
+        super().save_pretrained(save_directory, *args, **kwargs)
+        # The saved directory must work with `predict`, which needs the tokenizer.
+        if self._tokenizer is not None or self.config._name_or_path:
+            self.tokenizer.save_pretrained(save_directory)
+
+    @property
+    def tokenizer(self):
+        """Tokenizer used by `predict`, loaded from the checkpoint on first use."""
+        if self._tokenizer is None:
+            self._tokenizer = AutoTokenizer.from_pretrained(
+                self.config._name_or_path, **self._tokenizer_kwargs
+            )
+        return self._tokenizer
+
+    @tokenizer.setter
+    def tokenizer(self, value):
+        self._tokenizer = value
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        **kwargs,
+    ) -> PiiMaskingOutput:
+        last_hidden_state = self.backbone(
+            input_ids=input_ids, attention_mask=attention_mask, **kwargs
+        ).last_hidden_state                                        # [B, T, 1024]
+        h = last_hidden_state.to(self.token_cls_head.weight.dtype)  # heads are fp32
+
+        if attention_mask is None:
+            pooled = h.mean(dim=1)
+        else:
+            mask = attention_mask.to(h.dtype).unsqueeze(-1)
+            pooled = (h * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1.0)
+
+        return PiiMaskingOutput(
+            logits=self.token_cls_head(h),
+            sensitivity_logits=self.sensitivity_head(pooled).squeeze(-1),
+            last_hidden_state=last_hidden_state,
         )
 
     @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(
+    def predict(
+        self,
+        text: str,
+        tokenizer=None,
+        max_length: int | None = None,
+    ) -> tuple[list[PredictedSpan], float]:
+        """Predicted PII spans, and the document sensitivity in [0, 1].
+
+        Text past `max_length` tokens is truncated, so spans there are not
+        reported. An empty document has no spans and sensitivity 0.0.
+        """
+        tokenizer = tokenizer if tokenizer is not None else self.tokenizer
+        enc = tokenizer(text, return_offsets_mapping=True, return_tensors="pt",
+                        truncation=True,
+                        max_length=max_length or self.config.max_seq_len)
+        if enc["input_ids"].shape[1] == 0:  # empty document, nothing to encode
+            return [], 0.0
+        out = self(
             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
+        spans = self.viterbi.decode(out.logits[0].float().cpu(), offsets, text=text)
+        return spans, float(out.sensitivity_logits.float().sigmoid())
 
-    def mask(self, text: str) -> str:
-        spans, _ = self(text)
+    def mask(self, text: str, placeholder: str = "[{label}]", **kwargs) -> str:
+        """The text with every predicted span replaced by e.g. [PRIVATE_EMAIL]."""
+        spans, _ = self.predict(text, **kwargs)
         for s in sorted(spans, key=lambda s: -s.start):
-            text = text[:s.start] + f"[{s.label.upper()}]" + text[s.end:]
+            filled = placeholder.format(label=s.label.upper())
+            text = text[:s.start] + filled + 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))

modeling_pplx_qwen3.py

@@ -0,0 +1,94 @@
+# Vendored from perplexity-ai/pplx-embed-v1-0.6b @ 2c4d510dd4a732063c31a0f70193e35067b51fd8
+# (modeling.py). Kept in this repo so loading needs no code from another repo.
+import inspect
+from typing import Callable
+import torch
+from transformers import Qwen3Model
+from transformers.cache_utils import Cache
+from transformers.masking_utils import create_causal_mask
+from transformers.modeling_outputs import BaseModelOutputWithPooling
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs
+from .configuration_pplx_qwen3 import PPLXQwen3Config
+
+# The transformers `create_causal_mask` signature has shifted over releases:
+#   * <= 5.1: kwarg is `input_embeds`, `cache_position` is required positional
+#   * 5.2 - 5.5: renamed to `inputs_embeds`, `cache_position` still required
+#   * 5.6 - 5.8: `cache_position` has a default (kept for BC)
+#   * >= 5.9: `cache_position` removed entirely
+# Detect once at import time which names this transformers exposes.
+_CCM_PARAMS = inspect.signature(create_causal_mask).parameters
+_CCM_EMBEDS_KEY = "inputs_embeds" if "inputs_embeds" in _CCM_PARAMS else "input_embeds"
+_CCM_ACCEPTS_CACHE_POSITION = "cache_position" in _CCM_PARAMS
+
+
+# From modeling_t5gemma.py
+def bidirectional_mask_function(attention_mask: torch.Tensor | None) -> Callable:
+    """
+    This creates bidirectional attention mask.
+    """
+
+    def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool:
+        if attention_mask is None:
+            return torch.ones((), dtype=torch.bool)
+        return attention_mask[batch_idx, kv_idx].to(torch.bool)
+
+    return inner_mask
+
+
+class PPLXQwen3Model(Qwen3Model):
+    _supports_flash_attn = True
+    _supports_sdpa = True
+
+    config_class = PPLXQwen3Config
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.post_init()
+
+    def post_init(self):
+        super().post_init()
+        # Override to set all layers to non-causal attention. This'll work with attn_implementation="flash_attention_2" or "sdpa"
+        for layer in self.layers:
+            layer.self_attn.is_causal = False
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        use_cache: bool | None = None,
+        cache_position: torch.LongTensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPooling:
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+            input_ids = None
+
+        mask_kwargs = {
+            "config": self.config,
+            _CCM_EMBEDS_KEY: inputs_embeds,
+            "attention_mask": attention_mask,
+            "past_key_values": None,
+            "position_ids": position_ids,
+            "or_mask_function": bidirectional_mask_function(attention_mask),
+        }
+        if _CCM_ACCEPTS_CACHE_POSITION:
+            mask_kwargs["cache_position"] = torch.arange(
+                inputs_embeds.shape[1], device=inputs_embeds.device, dtype=torch.long
+            )
+        attention_mask = {"full_attention": create_causal_mask(**mask_kwargs)}
+
+        outputs = super().forward(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            **kwargs,
+        )
+        return outputs