Prepare 1.1.0: batch API (predict_many, scrub_many), opt-in float16

CHANGELOG.md

@@ -8,6 +8,27 @@ Semver for this island:
 
 Accuracy is 841-dev, union at 0.95, 354 gold spans. A version that changes those numbers must update `hybrid.json` `eval` and the tables below.
 
+## 1.1.0
+
+Same weights, rules, threshold and default outputs; `hybrid.json` `eval` is unchanged. New batch API and opt-in fp16.
+
+- **`Nergal.predict_many(texts)` / `scrub_many(texts)`:** windows from up to 64 texts are sorted by token length and packed into batches of at most 32,768 padded tokens and 128 rows. `predict` / `scrub` are now the one-text case of these.
+- **`dtype='float16'`** on `Nergal(...)` / `from_pretrained(...)` (CLI `--dtype`): casts the float32 weights at load time. Needs CUDA or MPS. The default stays `float32`, and `model.safetensors` is unchanged.
+- **Faster window sizing:** `Encoding.count` adds up cached unit pieces instead of re-tokenizing inside the window search. It yields the same windows, because `encode()` still rejects any unit whose pieces change with context.
+- **Default device:** `from_pretrained` now tries CUDA, then MPS, then CPU. 1.0.x used CPU even on CUDA machines.
+
+Throughput on one RTX 4090 (13.88M chars of FineWeb-2, kchar/s): 1.0.3-style per-document batches 11.0 (23.0 with 3 processes); `predict_many` float32 21.3 (28.1 with 2); `predict_many` float16 39.4 (79.6 with 3). The float16 path is limited by CPU-side tokenization, so run 2–3 processes per GPU.
+
+Equivalence: on the 1,685 labelled dev rows, float32 `predict_many` gives 0 span changes at 0.95 against the cached model spans of the published weights (841-dev max score change 2.6e-5). Float16 adds 2 spans on gold (1 on 841-dev, already covered by the rules) and removes none; 841-dev union numbers are identical. On 6,182 FineWeb-2 documents (1,236 with MPS and 4,946 with CUDA float32 references), CUDA float16 gave 0 span changes (max score change 0.012).
+
+| Version | Whole /354 | Residual | Rules FP | Union FP | Char P | Char R | What changed |
+|---|---:|---:|---:|---:|---:|---:|---|
+| 1.0.0 | 323 | 25 | 133 | 133 | 97.76% | 95.95% | First Hub snapshot |
+| 1.0.1 | 323 | 25 | 98 | 123 | 97.93% | 95.95% | Prefix-only glued-email trim |
+| 1.0.2 | 324 | 24 | 98 | 123 | 97.93% | 96.12% | Labelled country-area phone fix |
+| 1.0.3 | 324 | 24 | 98 | 123 | 97.93% | 96.12% | Label-note, e-Delivery and registry rules; placeholder and card fixes |
+| 1.1.0 | 324 | 24 | 98 | 123 | 97.93% | 96.12% | Batch API (`predict_many`, `scrub_many`), opt-in float16 |
+
 ## 1.0.3
 
 Same weights, threshold, and API. Rules SHA `f32d5c54…`.

README.md

@@ -13,7 +13,7 @@ tags:
 - hybrid
 ---
 
-# NERGAL 1.0.3
+# NERGAL 1.1.0
 
 **Named Entity Recognition with Grounded Additive Labels**
 
@@ -23,7 +23,7 @@ SlayerLab hybrid PII cleaner for Polish. Not a chat model. Not a drop-in `pipeli
 
 Python rules do the identifiers they can prove. A transformer NER head adds phone and other PII spans the regex misses. The cleaner **unions** the two on the original text, then replaces hits with `[Telefon]` or `[PII]`.
 
-- **Version:** `1.0.3` (`hybrid.json`, `CHANGELOG.md`)
+- **Version:** `1.1.0` (`hybrid.json`, `CHANGELOG.md`)
 - **Ground:** `scrub_pii` regex (SHA256 `f32d5c54…`)
 - **Additive labels:** XLM-RoBERTa-large token classifier, BIO tags `phone` / `pii`, threshold 0.95
 - **This snapshot:** seed `202609160`, **epoch 5** of a seven-epoch schedule
@@ -61,9 +61,9 @@ NERGAL is not designed to remove:
 
 These are intended exclusions; false positives can still mask some of this content.
 
-### Known gaps in 1.0.3
+### Known gaps in 1.1.0
 
-Unlabelled phones and identifiers, unusual formatting and damaged text can escape detection. **VINs and obfuscated emails** (such as `name (at) domain.pl`) are approved annotation targets, but that approval alone does not establish reliable support in the released 1.0.3 model. Do not rely on it to remove them consistently.
+Unlabelled phones and identifiers, unusual formatting and damaged text can escape detection. **VINs and obfuscated emails** (such as `name (at) domain.pl`) are approved annotation targets, but that approval alone does not establish reliable support in the released 1.1.0 model. Do not rely on it to remove them consistently.
 
 ## Versions
 
@@ -74,7 +74,8 @@ Unlabelled phones and identifiers, unusual formatting and damaged text can escap
 | 1.0.0 | 323 | 25 | 133 | 133 | 97.76% | 95.95% | First Hub snapshot |
 | 1.0.1 | 323 | 25 | 98 | 123 | 97.93% | 95.95% | Prefix-only glued-email trim |
 | 1.0.2 | 324 | 24 | 98 | 123 | 97.93% | 96.12% | Labelled country-area phone fix |
-| **1.0.3** | **324** | **24** | **98** | **123** | **97.93%** | **96.12%** | Label-note, e-Delivery and registry rules; placeholder and card fixes |
+| 1.0.3 | 324 | 24 | 98 | 123 | 97.93% | 96.12% | Label-note, e-Delivery and registry rules; placeholder and card fixes |
+| **1.1.0** | **324** | **24** | **98** | **123** | **97.93%** | **96.12%** | Batch API (`predict_many`, `scrub_many`), opt-in float16 |
 
 ## 841-dev
 
@@ -116,7 +117,7 @@ Same 841-dev split, threshold **0.95**. **Naked** is the transformer alone. **
 | Historical GLiNER email12 | naked | 249 | 70 | 10 | 99.81% | 85.39% |
 | Historical GLiNER email12 | ∪ regex | 290 | 50 | 108 | 98.13% | 93.76% |
 | XLM-R epoch 5 | naked | 298 | 42 | 57 | 98.96% | 89.38% |
-| **NERGAL 1.0.3** | **∪ regex** | 324 | 24 | 123 | 97.93% | 96.12% |
+| **NERGAL 1.1.0** | **∪ regex** | 324 | 24 | 123 | 97.93% | 96.12% |
 
 Zero-shot [GLiNER 2.5](https://huggingface.co/fastino/gliner2.5-multi-v1) is not competitive here, especially on non-phone PII (7/185 whole vs 160 naked / 179 union). Fine-tuned historical GLiNER is the precise naked baseline (10 false characters) and still trails XLM-R on coverage. NERGAL is XLM-R epoch 5 plus the regex: 145/169 phone, 179/185 other PII. Exact-span precision 86.34%, recall 89.27%, F1 87.78%.
 
@@ -151,6 +152,8 @@ nergal = Nergal.from_pretrained(root, local_files_only=True)
 masked, counts = nergal.scrub(text)
 ```
 
-`hybrid.json` records version `1.0.3`, threshold 0.95, gap ids `250002` / `250003`, the 841-dev `eval` block, and two weight hashes: `model_safetensors_sha256` for the published file and `source_checkpoint_sha256` for the training checkpoint it was packed from. `test_nergal.py` is synthetic (no corpus text). From this snapshot: `python -m unittest test_nergal`.
+For many texts, `nergal.scrub_many(texts)` (or `predict_many` for the raw model spans) batches windows across texts: about 2× the throughput of calling `scrub` in a loop on a CUDA GPU. `from_pretrained(..., dtype="float16")` casts the weights at load time on CUDA or MPS: another 1.9× on an RTX 4090, and 841-dev union numbers are unchanged, but scores are not bit-identical to float32 (0.95 spans can differ in rare cases). For a corpus, run 2–3 processes per GPU, because float16 inference is limited by CPU-side tokenization. Measurements are in `CHANGELOG.md`.
+
+`hybrid.json` records version `1.1.0`, threshold 0.95, gap ids `250002` / `250003`, the 841-dev `eval` block, and two weight hashes: `model_safetensors_sha256` for the published file and `source_checkpoint_sha256` for the training checkpoint it was packed from. `test_nergal.py` is synthetic (no corpus text). From this snapshot: `python -m unittest test_nergal`.
 
 Base weights: [`FacebookAI/xlm-roberta-large`](https://huggingface.co/FacebookAI/xlm-roberta-large) revision `c23d21b0620b635a76227c604d44e43a9f0ee389` (MIT).

hybrid.json

@@ -1,7 +1,7 @@
 {
   "full_name": "Named Entity Recognition with Grounded Additive Labels",
   "hub_id": "SlayerLab/NERGAL",
-  "version": "1.0.3",
+  "version": "1.1.0",
   "mode": "rules_union",
   "epoch": 5,
   "seed": 202609160,

nergal.py

@@ -17,13 +17,14 @@ import scrub_pii
 from scrub_pii import PHONE_TAG, PII_TAG
 
 HUB_ID = 'SlayerLab/NERGAL'
-VERSION = '1.0.3'
+VERSION = '1.1.0'
 GAPS = ['[PII_SPACE]', '[PII_BREAK]']
 GAP_IDS = [250002, 250003]
 BIO_LABELS = ['O', 'B-phone', 'I-phone', 'B-pii', 'I-pii']
 LABELS = ['phone', 'pii']
 THRESHOLD = 0.95
 RULES_SHA = 'f32d5c5452fc47178e109d4bc248a0d8234ea6e59e8cf79407f4eb8451581d67'
+DTYPES = ('float32', 'float16')
 
 
 def sha(path):
@@ -157,8 +158,9 @@ class Encoding:
         return encoded, [first[j] for j in range(len(words))]
 
     def count(self, units):
-        encoded, _ = self.encode([u.model for u in units])
-        return len(encoded['input_ids'][0])
+        """Special tokens + cached unit pieces. Exact because encode() rejects any unit whose pieces change with
+        context; no re-tokenization inside the windows() binary search."""
+        return self.tokenizer.num_special_tokens_to_add() + sum(len(self.pieces(u.model)) for u in units)
 
     def prepare(self, text):
         units = unitize(text, self.pieces, self.tokenizer.unk_token)
@@ -245,9 +247,14 @@ def _load_rules_module(asset):
 
 
 class Nergal:
-    def __init__(self, asset, device='cpu'):
+    def __init__(self, asset, device='cpu', dtype='float32'):
         import torch
         from transformers import AutoModelForTokenClassification, AutoTokenizer
+        if dtype not in DTYPES:
+            raise ValueError(f'dtype must be one of {DTYPES}')
+        if dtype == 'float16' and torch.device(device).type == 'cpu':
+            raise ValueError('float16 needs a cuda or mps device')
+        self.dtype = dtype
         self.device = device
         self._torch = torch
         asset = Path(asset)
@@ -263,47 +270,74 @@ class Nergal:
         self.model = AutoModelForTokenClassification.from_pretrained(str(asset), local_files_only=True)
         self.encoding = Encoding(tokenizer)
         self.threshold = THRESHOLD
-        self.model.to(device).eval()
+        self.model.to(device, dtype=getattr(torch, dtype)).eval()
 
     @classmethod
-    def from_pretrained(cls, source=HUB_ID, *, device=None, local_files_only=False):
+    def from_pretrained(cls, source=HUB_ID, *, device=None, dtype='float32', local_files_only=False):
         import torch
         if device is None:
-            device = 'mps' if torch.backends.mps.is_available() else 'cpu'
-        return cls(_resolve(source, local_files_only=local_files_only), device=device)
+            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)
 
     def predict(self, text):
-        torch = self._torch
-        units, chunks = self.encoding.prepare(text)
-        if not units:
-            return []
-        sums, counts = torch.zeros(len(units), 5), torch.zeros(len(units), 1)
-        with torch.inference_mode():
+        return self.predict_many([text])[0]
+
+    def predict_many(self, texts, *, batch_tokens=32768, max_batch=128, group=64):
+        """Model spans for each text. Windows of up to `group` texts are sorted by token length and packed into
+        batches of at most `max_batch` rows and `batch_tokens` padded tokens; per text, units, windows, averaging
+        and decode are the same as one window at a time."""
+        if min(batch_tokens, max_batch, group) < 1:
+            raise ValueError('batch_tokens, max_batch and group must be positive')
+        texts = list(texts)
+        result = []
+        for i in range(0, len(texts), group):
+            result += self._predict_group(texts[i:i + group], batch_tokens, max_batch)
+        return result
+
+    def _predict_group(self, texts, batch_tokens, max_batch):
+        torch, encoding = self._torch, self.encoding
+        docs, items = [], []
+        for d, text in enumerate(texts):
+            units, chunks = encoding.prepare(text)
+            docs.append((units, torch.zeros(len(units), 5), torch.zeros(len(units), 1)))
             for window in chunks:
-                a, b = window['start'], window['end']
-                words = [u.model for u in units[a:b]]
-                encoded, first = self.encoding.encode(words)
-                batch = self.encoding.tokenizer.pad(
-                    [{k: v[0] for k, v in encoded.items()}], padding=True, return_tensors='pt',
-                )
-                batch = {k: v.to(self.device) if torch.is_tensor(v) else v for k, v in batch.items()}
+                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]
+                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:
                     raise ValueError('Batch exceeds encoder limit')
-                logits = self.model(**batch).logits
-                sums[a:b] += logits[0, first].float().cpu()
-                counts[a:b] += 1
-        if (counts == 0).any():
-            raise ValueError('Missing inference units')
-        return decode_bio(units, (sums / counts).tolist())
+                logits = self.model(**{k: v.to(self.device) for k, v in batch.items()}).logits.float().cpu()
+                for row, (_, d, window, _, first) in enumerate(part):
+                    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():
+                raise ValueError('Missing inference units')
+            result.append(decode_bio(units, (sums / counts).tolist()) if units else [])
+        return result
 
     def rule_spans(self, text):
         return spans_from(self._scrub, text)
 
     def scrub(self, text):
-        if not text:
-            return text, {'phone': 0, 'pii': 0, 'rules_placeholder_chars': 0,
-                          'union_placeholder_chars': 0, 'model_extra_spans': 0}
-        return scrub_spans(text, self.rule_spans(text), self.predict(text), threshold=self.threshold)
+        return self.scrub_many([text])[0]
+
+    def scrub_many(self, texts):
+        texts = list(texts)
+        return [scrub_spans(text, self.rule_spans(text), model, threshold=self.threshold)
+                for text, model in zip(texts, self.predict_many(texts), strict=True)]
 
 
 def main(argv=None):
@@ -312,9 +346,11 @@ def main(argv=None):
     parser = argparse.ArgumentParser(description='NERGAL hybrid PII cleaner')
     parser.add_argument('--repo', default=HUB_ID)
     parser.add_argument('--device', default=None)
+    parser.add_argument('--dtype', default='float32', choices=DTYPES)
     parser.add_argument('--local', action='store_true')
     args = parser.parse_args(argv)
-    nergal = Nergal.from_pretrained(args.repo, device=args.device, local_files_only=args.local)
+    nergal = Nergal.from_pretrained(args.repo, device=args.device, dtype=args.dtype,
+                                     local_files_only=args.local)
     text = sys.stdin.read()
     masked, counts = nergal.scrub(text)
     sys.stdout.write(masked)

test_nergal.py

@@ -13,7 +13,7 @@ class NergalTests(unittest.TestCase):
         from nergal import GAP_IDS, GAPS, HUB_ID, RULES_SHA as PINNED, THRESHOLD, VERSION
         card = json.loads((HERE / 'hybrid.json').read_text())
         self.assertEqual(HUB_ID, 'SlayerLab/NERGAL')
-        self.assertEqual(VERSION, '1.0.3')
+        self.assertEqual(VERSION, '1.1.0')
         self.assertEqual(card['version'], VERSION)
         self.assertEqual(card['eval']['union_fp'], 123)
         self.assertEqual(card['eval']['rules_fp'], 98)
@@ -35,6 +35,26 @@ class NergalTests(unittest.TestCase):
         self.assertEqual(len(first), len(words))
         self.assertEqual([encoded.word_ids(0)[i] for i in first], [0, 1, 2])
 
+    def test_window_token_count_matches_the_encoded_window(self):
+        from transformers import AutoTokenizer
+        from nergal import Encoding
+        tokenizer = AutoTokenizer.from_pretrained(str(HERE), local_files_only=True, fix_mistral_regex=False)
+        encoding = Encoding(tokenizer)
+        text = ' '.join(f'Zdanie {i}: tel. 22 123 45 67,\nNIP 1234567802.' for i in range(120))
+        units, chunks = encoding.prepare(text)
+        self.assertGreater(len(chunks), 1)
+        for w in chunks:
+            encoded, _ = encoding.encode([u.model for u in units[w['start']:w['end']]])
+            self.assertEqual(w['tokens'], len(encoded['input_ids'][0]))
+            self.assertLessEqual(w['tokens'], 512)
+
+    def test_float16_is_opt_in_and_needs_an_accelerator(self):
+        from nergal import Nergal
+        with self.assertRaises(ValueError):
+            Nergal(HERE, device='cpu', dtype='float16')
+        with self.assertRaises(ValueError):
+            Nergal(HERE, dtype='bfloat16')
+
     def test_existing_placeholders_do_not_switch_the_rules_off(self):
         from nergal import rules
         text = 'Kontakt [Telefon], NIP 1234567802.'  # invented, checksum-valid