Add an AutoModel custom code implementation

README.md

@@ -27,38 +27,37 @@ 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 loads the
+encoder implementation 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.
+
 ## Checkpoint layout
 
 `model.safetensors` holds the fine-tuned backbone (bf16, `backbone.*`), both
@@ -69,5 +68,5 @@ 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
+referenced by `config.json`'s `backbone.auto_map`) ships with the
 [backbone repo](https://huggingface.co/perplexity-ai/pplx-embed-v1-0.6b).

config.json

@@ -1,4 +1,12 @@
 {
+  "architectures": [
+    "PiiMaskingModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "modeling_pii_masking.PiiMaskingConfig",
+    "AutoModel": "modeling_pii_masking.PiiMaskingModel"
+  },
+  "dtype": "bfloat16",
   "model_type": "pii_masking",
   "backbone": {
     "vocab_size": 151936,

example_usage.py → modeling_pii_masking.py

@@ -1,23 +1,28 @@
-#!/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 pulled from the public backbone
+repo perplexity-ai/pplx-embed-v1-0.6b, 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.dynamic_module_utils import get_class_from_dynamic_module
+from transformers.utils import ModelOutput
 
 BACKBONE_REPO = "perplexity-ai/pplx-embed-v1-0.6b"
 
@@ -66,24 +71,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 +106,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 +227,140 @@ 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 because the encoder config class lives in the
+    backbone repo rather than in transformers.
+    """
+
+    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)
+        # Encoder architecture from the public backbone repo. The weights come
+        # from this checkpoint's backbone.* tensors.
+        backbone_config = get_class_from_dynamic_module(
+            "configuration.PPLXQwen3Config", BACKBONE_REPO
+        ).from_dict(config.backbone)
+        backbone_config._attn_implementation = config._attn_implementation
+        self.backbone = get_class_from_dynamic_module(
+            "modeling.PPLXQwen3Model", BACKBONE_REPO
+        )(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.post_init()
 
-    @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)
+    @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)
+        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:
         h = self.backbone(
+            input_ids=input_ids, attention_mask=attention_mask, **kwargs
+        ).last_hidden_state.to(self.token_cls_head.weight.dtype)   # [B, T, 1024]
+
+        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=h,
+        )
+
+    @torch.no_grad()
+    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]."""
+        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))