2.0.0: opt-in FastPDN person names (names/, CC-BY-4.0)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

hybrid.json

@@ -46,5 +46,23 @@
   "backbone": {
     "repo": "FacebookAI/xlm-roberta-large",
     "revision": "c23d21b0620b635a76227c604d44e43a9f0ee389"
+  },
+  "names": {
+    "repo": "ArkadiuszPawlak/fastpdn-ner-polish-pii",
+    "revision": "636c57fe77b0afad8f065e7c76243dc051437035",
+    "license": "CC-BY-4.0",
+    "max_length": 512,
+    "stride": 128,
+    "dtype": "float32",
+    "decode": "argmax, union over windows, expanded to whole words",
+    "converted_from": {
+      "model.onnx": "11924bfd20eb478ee614e5e1184536825db4747753d341af7c04b940efa84645",
+      "convert.py": "a698bc993b253fc70d960aa14af37a5088afb79c77dae70557c7056ce394333e"
+    },
+    "sha256": {
+      "config.json": "3cef51cd7d184478adbdec9f68730db279386fe145cb32473a0622edea923578",
+      "model.safetensors": "2fa84ec6abd0c1b12ba0ba1548e89c9fca36fc58f03e66074db423b0e98ec8e6",
+      "tokenizer.json": "3a690c2d605076ad3901d946d4c0145fbfb19fef2f01370c7388430aa9f31edf"
+    }
   }
 }

names/NOTICE.md

@@ -0,0 +1,29 @@
+# names/: FastPDN NER — Polish PII
+
+- Source: [ArkadiuszPawlak/fastpdn-ner-polish-pii](https://huggingface.co/ArkadiuszPawlak/fastpdn-ner-polish-pii),
+  revision `636c57fe77b0afad8f065e7c76243dc051437035`.
+- License: [CC-BY-4.0](https://creativecommons.org/licenses/by/4.0/).
+- Attribution: "FastPDN NER — Polish PII (ONNX)" by ArkadiuszPawlak, fine-tuned from
+  [clarin-pl/FastPDN](https://huggingface.co/clarin-pl/FastPDN) on filtered KPWr
+  ([clarin-pl/kpwr-ner](https://huggingface.co/datasets/clarin-pl/kpwr-ner)) plus LLM-synthetic data.
+  Upstream model card README.md sha256 `9ca776b1157e5dedfb9d404ee17a584ebfc80f4dc426beb68f69200a52522c6f`.
+
+## Changes
+
+- `config.json` and `tokenizer.json` unchanged.
+- `model.safetensors` is upstream `model.onnx` (FP32, sha256
+  `11924bfd20eb478ee614e5e1184536825db4747753d341af7c04b940efa84645`) converted by `convert.py` (sha256
+  `a698bc993b253fc70d960aa14af37a5088afb79c77dae70557c7056ce394333e`): 199 tensors copied, 73 linear weights
+  transposed from ONNX [in, out] to torch [out, in], no value changed.
+- Review Gate 14, on 1,000 documents (1,714 windows, 614,466 tokens): torch FP32 against ONNX Runtime FP32, max
+  absolute logit difference 3.4e-4 (within 1e-3), every token's argmax identical, every person range identical.
+- NERGAL uses only the PERSON, PERSON_F and PERSON_L tags, masked as `[PERSON]`.
+
+## Files
+
+| File | sha256 |
+|---|---|
+| `config.json` | `3cef51cd7d184478adbdec9f68730db279386fe145cb32473a0622edea923578` |
+| `model.safetensors` | `2fa84ec6abd0c1b12ba0ba1548e89c9fca36fc58f03e66074db423b0e98ec8e6` |
+| `tokenizer.json` | `3a690c2d605076ad3901d946d4c0145fbfb19fef2f01370c7388430aa9f31edf` |
+| `NOTICE.md` | this file |

names/config.json

@@ -0,0 +1,60 @@
+{
+  "add_cross_attention": false,
+  "architectures": [
+    "BertForTokenClassification"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "bos_token_id": null,
+  "classifier_dropout": 0.1,
+  "dtype": "float32",
+  "eos_token_id": null,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "id2label": {
+    "0": "O",
+    "1": "B-PERSON",
+    "2": "I-PERSON",
+    "3": "B-PERSON_F",
+    "4": "I-PERSON_F",
+    "5": "B-PERSON_L",
+    "6": "I-PERSON_L",
+    "7": "B-ORG",
+    "8": "I-ORG",
+    "9": "B-STREET",
+    "10": "I-STREET",
+    "11": "B-CITY",
+    "12": "I-CITY"
+  },
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "is_decoder": false,
+  "label2id": {
+    "B-CITY": 11,
+    "B-ORG": 7,
+    "B-PERSON": 1,
+    "B-PERSON_F": 3,
+    "B-PERSON_L": 5,
+    "B-STREET": 9,
+    "I-CITY": 12,
+    "I-ORG": 8,
+    "I-PERSON": 2,
+    "I-PERSON_F": 4,
+    "I-PERSON_L": 6,
+    "I-STREET": 10,
+    "O": 0
+  },
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 514,
+  "model_type": "bert",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 1,
+  "position_embedding_type": "absolute",
+  "tie_word_embeddings": true,
+  "tokenizer_class": "HerbertTokenizerFast",
+  "transformers_version": "5.12.1",
+  "type_vocab_size": 2,
+  "use_cache": false,
+  "vocab_size": 50000
+}

names/model.safetensors

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

names/tokenizer.json

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

nergal.py

@@ -258,12 +258,13 @@ def scrub_spans(text, rule_spans, model_spans, *, threshold=THRESHOLD, module=sc
                     'model_extra_spans': extra}
 
 
-def _resolve(source, *, local_files_only):
+def _resolve(source, *, local_files_only, names=False):
     path = Path(source)
     if path.is_dir():
         return path
     from huggingface_hub import snapshot_download
-    return Path(snapshot_download(source, local_files_only=local_files_only))
+    return Path(snapshot_download(source, local_files_only=local_files_only,
+                                  ignore_patterns=None if names else ['names/*']))  # names/ is ≈490 MB
 
 
 def _load_rules_module(asset):
@@ -279,8 +280,74 @@ def _load_rules_module(asset):
     return scrub_pii
 
 
+_FEEDS = (('input_ids', 'ids'), ('attention_mask', 'attention_mask'), ('token_type_ids', 'type_ids'))
+
+
+def pack(lengths, batch_tokens, max_batch):
+    """Indices sorted by length, cut into batches of at most max_batch rows and batch_tokens padded tokens."""
+    order = sorted(range(len(lengths)), key=lengths.__getitem__)
+    out, i = [], 0
+    while i < len(order):
+        j = i + 1  # ascending, so order[j] sets the padded width of order[i:j + 1]
+        while j < len(order) and j - i < max_batch and (j - i + 1) * lengths[order[j]] <= batch_tokens:
+            j += 1
+        out.append(order[i:j])
+        i = j
+    return out
+
+
+class Names:
+    """FastPDN-PII person names (names/, CC-BY-4.0, see names/NOTICE.md): the FP32 ONNX converted to safetensors, run
+    in torch on NERGAL's device; argmax per token, windows packed like predict_many."""
+
+    def __init__(self, asset, card, device='cpu'):
+        import torch
+        from safetensors.torch import load_file
+        from tokenizers import Tokenizer
+        from transformers import BertConfig, BertForTokenClassification
+        root = Path(asset) / 'names'
+        for name, digest in card['sha256'].items():
+            if sha(root / name) != digest:
+                raise ValueError(f'Unexpected names/{name} sha256')
+        config = BertConfig.from_pretrained(str(root))
+        labels = config.id2label
+        self.person_ids = {i for i, label in labels.items() if label[2:] in PERSON_LABELS}
+        if sorted(labels[i] for i in self.person_ids) != sorted(f'{p}-{l}' for p in 'BI' for l in PERSON_LABELS):
+            raise ValueError('names/config.json does not label B/I person tags')
+        self.tokenizer = Tokenizer.from_file(str(root / 'tokenizer.json'))
+        self.tokenizer.enable_truncation(max_length=card['max_length'], stride=card['stride'])
+        self.tokenizer.no_padding()
+        self.pad = {'input_ids': config.pad_token_id, 'attention_mask': 0, 'token_type_ids': 0}
+        self.model = BertForTokenClassification(config)
+        self.model.load_state_dict(load_file(root / 'model.safetensors'), strict=True)
+        self.model.to(device, dtype=getattr(torch, card['dtype'])).eval()
+        self.device, self._torch = device, torch
+
+    def spans(self, text):
+        return self.spans_many([text])[0]
+
+    def spans_many(self, texts, *, batch_tokens=32768, max_batch=128):
+        torch, texts = self._torch, list(texts)
+        items = [(d, w) for d, text in enumerate(texts)
+                 for enc in [self.tokenizer.encode(text)] for w in [enc, *enc.overflowing]]
+        ranges = [[] for _ in texts]
+        with torch.inference_mode():
+            for part in pack([len(w.ids) for _, w in items], batch_tokens, max_batch):
+                rows = [items[i][1] for i in part]
+                width = max(len(w.ids) for w in rows)
+                feed = {name: torch.tensor([getattr(w, attr) + [self.pad[name]] * (width - len(w.ids)) for w in rows],
+                                           device=self.device) for name, attr in _FEEDS}
+                logits = self.model(**feed).logits.float().cpu()
+                for row, i in enumerate(part):
+                    d, w = items[i]
+                    tags = logits[row, :len(w.ids)].argmax(-1).tolist()
+                    ranges[d] += [(a, b) for tag, special, (a, b) in zip(tags, w.special_tokens_mask, w.offsets, strict=True)
+                                  if tag in self.person_ids and not special and a < b]
+        return [person_spans(text, r) for text, r in zip(texts, ranges, strict=True)]
+
+
 class Nergal:
-    def __init__(self, asset, device='cpu', dtype='float32'):
+    def __init__(self, asset, device='cpu', dtype='float32', names=False):
         import torch
         from transformers import AutoModelForTokenClassification, AutoTokenizer
         if dtype not in DTYPES:
@@ -295,6 +362,8 @@ class Nergal:
         card = json.loads((asset / 'hybrid.json').read_text())
         if card['gap_ids'] != GAP_IDS or card['threshold'] != THRESHOLD:
             raise ValueError('hybrid.json does not match this NERGAL snapshot')
+        if names and 'names' not in card:
+            raise ValueError('This NERGAL snapshot has no names model')
         tokenizer = AutoTokenizer.from_pretrained(
             str(asset), local_files_only=True, use_fast=True, fix_mistral_regex=False,
         )
@@ -304,13 +373,15 @@ class Nergal:
         self.encoding = Encoding(tokenizer)
         self.threshold = THRESHOLD
         self.model.to(device, dtype=getattr(torch, dtype)).eval()
+        self.names = Names(asset, card['names'], device) if names else None  # dtype from the card, not NERGAL's
 
     @classmethod
-    def from_pretrained(cls, source=HUB_ID, *, device=None, dtype='float32', local_files_only=False):
+    def from_pretrained(cls, source=HUB_ID, *, device=None, dtype='float32', local_files_only=False, names=False):
         import torch
         if device is None:
             device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
-        return cls(_resolve(source, local_files_only=local_files_only), device=device, dtype=dtype)
+        return cls(_resolve(source, local_files_only=local_files_only, names=names), device=device, dtype=dtype,
+                   names=names)
 
     def predict(self, text):
         return self.predict_many([text])[0]
@@ -336,14 +407,9 @@ class Nergal:
             for window in chunks:
                 encoded, first = encoding.encode([u.model for u in units[window['start']:window['end']]])
                 items.append((len(encoded['input_ids'][0]), d, window, encoded, first))
-        items.sort(key=lambda item: item[0])
         with torch.inference_mode():
-            i = 0
-            while i < len(items):
-                j = i + 1  # sorted ascending, so items[j] sets the padded width of items[i:j + 1]
-                while j < len(items) and j - i < max_batch and (j - i + 1) * items[j][0] <= batch_tokens:
-                    j += 1
-                part = items[i:j]
+            for part in pack([n for n, *_ in items], batch_tokens, max_batch):
+                part = [items[i] for i in part]
                 batch = encoding.tokenizer.pad([{k: v[0] for k, v in e.items()} for _, _, _, e, _ in part],
                                                padding=True, return_tensors='pt')
                 if batch['input_ids'].shape[1] > 512:
@@ -353,7 +419,6 @@ class Nergal:
                     a, b = window['start'], window['end']
                     docs[d][1][a:b] += logits[row, first]
                     docs[d][2][a:b] += 1
-                i = j
         result = []
         for units, sums, counts in docs:
             if (counts == 0).any():
@@ -369,8 +434,10 @@ class Nergal:
 
     def scrub_many(self, texts):
         texts = list(texts)
-        return [scrub_spans(text, self.rule_spans(text), model, threshold=self.threshold, module=self._scrub)
-                for text, model in zip(texts, self.predict_many(texts), strict=True)]
+        name_spans = self.names.spans_many(texts) if self.names else [()] * len(texts)
+        return [scrub_spans(text, self.rule_spans(text), model, threshold=self.threshold, module=self._scrub,
+                            name_spans=spans)
+                for text, model, spans in zip(texts, self.predict_many(texts), name_spans, strict=True)]
 
 
 def main(argv=None):
@@ -381,9 +448,10 @@ def main(argv=None):
     parser.add_argument('--device', default=None)
     parser.add_argument('--dtype', default='float32', choices=DTYPES)
     parser.add_argument('--local', action='store_true')
+    parser.add_argument('--names', action='store_true')
     args = parser.parse_args(argv)
     nergal = Nergal.from_pretrained(args.repo, device=args.device, dtype=args.dtype,
-                                     local_files_only=args.local)
+                                     local_files_only=args.local, names=args.names)
     text = sys.stdin.read()
     masked, counts = nergal.scrub(text)
     sys.stdout.write(masked)

test_nergal.py

@@ -1,6 +1,8 @@
 """Synthetic NERGAL tests. Invented strings only; no corpus text or real identifiers."""
 import hashlib
 import json
+import shutil
+import tempfile
 import unittest
 from pathlib import Path
 
@@ -172,6 +174,68 @@ class NergalTests(unittest.TestCase):
         self.assertEqual(scrub_spans(text, [], [span(text, '112')])[0], text)
         self.assertEqual(model_keep(text, [span(text, '112', score=0.9)], threshold=0.95), [])
 
+    def _names(self):
+        from nergal import Names
+        return Names(HERE, json.loads((HERE / 'hybrid.json').read_text())['names'])
+
+    def test_names_mask_an_invented_person_and_nothing_else(self):
+        names = self._names()
+        text = 'Wniosek złożyła Anna Nowakowska z Radomia.'
+        self.assertIn('Anna Nowakowska', [text[s['start']:s['end']] for s in names.spans(text)])
+        self.assertEqual(names.spans('Zdanie bez nazwisk o pogodzie.'), [])  # special tokens are never a person
+        self.assertEqual(names.spans(''), [])
+        text = 'Ala ma kota.'  # a first name alone is a person (names policy)
+        self.assertEqual([text[s['start']:s['end']] for s in names.spans(text)], ['Ala'])
+
+    def test_names_find_a_person_past_the_first_window_and_are_idempotent(self):
+        from nergal import apply_union
+        names = self._names()
+        text = 'Zdanie bez nazwisk o pogodzie. ' * 200 + 'Wniosek złożyła Anna Nowakowska.'
+        spans = names.spans(text)
+        self.assertEqual([text[s['start']:s['end']] for s in spans], ['Anna Nowakowska'])
+        masked, _, _ = apply_union(text, spans)
+        self.assertEqual(names.spans(masked), [])
+
+    def test_names_batched_equal_one_at_a_time(self):
+        names = self._names()
+        texts = ['Wniosek złożyła Anna Nowakowska z Radomia.', '',
+                 'Zdanie bez nazwisk o pogodzie. ' * 200 + 'Podpisał Tomasz Wrzos.', 'Ala ma kota.']
+        self.assertEqual(names.spans_many(texts, max_batch=2), [names.spans(t) for t in texts])
+
+    def test_names_are_opt_in_and_verified(self):
+        from nergal import Names
+        card = json.loads((HERE / 'hybrid.json').read_text())['names']
+        with self.assertRaises(ValueError):
+            Names(HERE, {**card, 'sha256': {**card['sha256'], 'tokenizer.json': '0' * 64}})
+
+    def test_a_snapshot_without_names_refuses_names(self):
+        from nergal import Nergal
+        card = json.loads((HERE / 'hybrid.json').read_text())
+        del card['names']
+        with tempfile.TemporaryDirectory() as tmp:
+            shutil.copy(HERE / 'scrub_pii.py', tmp)
+            (Path(tmp) / 'hybrid.json').write_text(json.dumps(card))
+            with self.assertRaisesRegex(ValueError, 'no names model'):
+                Nergal(tmp, names=True)  # raised before the tokenizer or model loads
+
+    def test_hub_download_skips_names_unless_asked(self):
+        from unittest import mock
+        import huggingface_hub
+        from nergal import _resolve
+        with mock.patch.object(huggingface_hub, 'snapshot_download', return_value=str(HERE)) as download:
+            _resolve('SlayerLab/NERGAL', local_files_only=True)
+            _resolve('SlayerLab/NERGAL', local_files_only=True, names=True)
+        self.assertEqual([c.kwargs['ignore_patterns'] for c in download.call_args_list], [['names/*'], None])
+
+    def test_pack_keeps_every_index_once_within_limits(self):
+        from nergal import pack
+        lengths = [5, 1, 9, 3, 3, 7]
+        parts = pack(lengths, batch_tokens=12, max_batch=2)
+        self.assertEqual(sorted(i for p in parts for i in p), list(range(6)))
+        for p in parts:
+            self.assertLessEqual(len(p), 2)
+            self.assertLessEqual(len(p) * max(lengths[i] for i in p), 12)
+
 
 if __name__ == '__main__':
     unittest.main()