Add hybrid PII scripts: nergal.py and frozen scrub_pii.py

nergal.py

@@ -0,0 +1,326 @@
+"""NERGAL hybrid PII cleaner: frozen regex ∪ windowed XLM-R BIO head.
+
+This file is the public PII island. It does not import the lab training stack and
+must not call embedding-extension. The packed tokenizer already has the gap ids.
+"""
+from __future__ import annotations
+
+import hashlib
+import json
+import math
+import re
+from dataclasses import dataclass
+from functools import lru_cache
+from pathlib import Path
+
+import scrub_pii
+from scrub_pii import PHONE_TAG, PII_TAG
+
+HUB_ID = 'SlayerLab/NERGAL'
+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 = '547c0428b0799bf051566d6ac489987eff27f09e1bda36a5665452fe155b3966'
+
+
+def sha(path):
+    with Path(path).open('rb') as stream:
+        return hashlib.file_digest(stream, 'sha256').hexdigest()
+
+
+def verify_rules(path=None):
+    digest = sha(path or scrub_pii.__file__)
+    if digest != RULES_SHA:
+        raise ValueError(f'Unexpected rules sha256 {digest}')
+    return digest
+
+
+@dataclass(frozen=True)
+class Unit:
+    model: str
+    start: int
+    end: int
+    gap: bool
+
+
+def unitize(text, encode=None, unk=None):
+    units = []
+    for match in re.finditer(r'\s+|\S', text):
+        raw = match.group()
+        gap = raw.isspace()
+        model = GAPS[int(any(c in raw for c in '\r\n\v\f\x85\u2028\u2029'))] if gap else raw
+        if not gap and encode is not None and not encode(model):
+            if not unk:
+                raise ValueError('Zero-piece unit without unknown token')
+            model = unk
+        units.append(Unit(model, match.start(), match.end(), gap))
+    return units
+
+
+def windows(units, count, *, max_units=384, limit=512):
+    overlap = 128
+    result, start = [], 0
+    while start < len(units):
+        lo, hi = start + 1, min(start + max_units, len(units))
+        while lo < hi:
+            mid = (lo + hi + 1) // 2
+            if count(units[start:mid]) <= limit:
+                lo = mid
+            else:
+                hi = mid - 1
+        end, size = lo, count(units[start:lo])
+        if size > limit or (end < len(units) and end - start <= overlap):
+            raise ValueError('Token budget cannot fit a progressing window')
+        width = 64
+        result.append({
+            'start': start, 'end': end, 'tokens': size,
+            'owner_start': start if not result else start + width - 1,
+            'owner_end': end if end == len(units) else end - width + 1,
+        })
+        if end == len(units):
+            break
+        start = end - overlap
+    return result
+
+
+def raw_span(units, a, b, label, score):
+    if not 0 <= a < b <= len(units) or label not in LABELS or not math.isfinite(score) or not 0 <= score <= 1:
+        raise ValueError('Invalid unit prediction')
+    while a < b and units[a].gap:
+        a += 1
+    while a < b and units[b - 1].gap:
+        b -= 1
+    if a == b:
+        return None
+    return {'start': units[a].start, 'end': units[b - 1].end, 'label': label, 'score': score}
+
+
+def decode_bio(units, logits):
+    if len(units) != len(logits) or any(len(v) != 5 or any(not math.isfinite(x) for x in v) for v in logits):
+        raise ValueError('Invalid BIO logits')
+    result, active, probabilities = [], None, []
+
+    def finish(end):
+        nonlocal active
+        if active is None:
+            return
+        start, label = active
+        span = raw_span(units, start, end, label, min(probabilities[start:end]))
+        if span is not None:
+            result.append(span)
+        active = None
+
+    for i, values in enumerate(logits):
+        tag = max(range(5), key=lambda j: values[j])
+        exponentials = [math.exp(x - max(values)) for x in values]
+        probabilities.append(exponentials[tag] / sum(exponentials))
+        label = LABELS[(tag - 1) // 2] if tag else None
+        if tag == 0 or tag in (1, 3) or active is None or active[1] != label:
+            finish(i)
+            active = (i, label) if tag else None
+    finish(len(units))
+    return result
+
+
+def decode(spans, threshold=THRESHOLD):
+    if any(not math.isfinite(s['score']) or not 0 <= s['score'] <= 1 for s in spans):
+        raise ValueError('Nonfinite/invalid confidence')
+    result = []
+    for span in sorted(spans, key=lambda s: (-s['score'], -(s['end'] - s['start']), s['start'], s['label'])):
+        if span['score'] >= threshold and not any(span['start'] < p['end'] and p['start'] < span['end'] for p in result):
+            result.append(span)
+    return sorted(result, key=lambda s: (s['start'], s['end'], s['label']))
+
+
+class Encoding:
+    def __init__(self, tokenizer):
+        self.tokenizer = tokenizer
+        tokenizer.model_max_length = 512
+        self.pieces = lru_cache(maxsize=16384)(lambda s: tuple(tokenizer.encode(s, add_special_tokens=False)))
+
+    def encode(self, words):
+        encoded = self.tokenizer(words, is_split_into_words=True, truncation=False, padding=False)
+        ids = encoded['input_ids'][0]
+        mapping = encoded.word_ids(0)
+        first, actual = {}, {}
+        for i, word in enumerate(mapping):
+            if word is not None:
+                first.setdefault(word, i)
+                actual.setdefault(word, []).append(ids[i])
+        if set(first) != set(range(len(words))):
+            raise ValueError('Tokenizer dropped a unit')
+        if any(tuple(actual[j]) != self.pieces(word) for j, word in enumerate(words)):
+            raise ValueError('Unit token IDs change with window context')
+        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])
+
+    def prepare(self, text):
+        units = unitize(text, self.pieces, self.tokenizer.unk_token)
+        return units, windows(units, self.count) if units else []
+
+
+def rules(text):
+    verify_rules()
+    result = []
+    scrub_pii.scrub_pii(text, spans=result)
+    if '[PII]' in text or '[Telefon]' in text:
+        return []
+    return sorted(({k: s[k] for k in ('start', 'end', 'label')} | {'score': 1.0} for s in result),
+                  key=lambda s: s['start'])
+
+
+def apply_union(text, spans):
+    labels = [None] * len(text)
+    for span in spans:
+        start, end, label = span['start'], span['end'], span['label']
+        if not 0 <= start < end <= len(text):
+            raise ValueError('Span outside text')
+        for i in range(start, end):
+            if labels[i] is None or label == 'phone':
+                labels[i] = label
+    out, chars, n_phone, n_pii, i = [], 0, 0, 0, 0
+    while i < len(text):
+        lab = labels[i]
+        if lab is None:
+            out.append(text[i])
+            i += 1
+            continue
+        j = i + 1
+        while j < len(text) and labels[j] == lab:
+            j += 1
+        tag = PHONE_TAG if lab == 'phone' else PII_TAG
+        out.append(tag)
+        chars += len(tag)
+        if lab == 'phone':
+            n_phone += 1
+        else:
+            n_pii += 1
+        i = j
+    return ''.join(out), chars, n_phone, n_pii
+
+
+def scrub_spans(text, rule_spans, model_spans, *, threshold=THRESHOLD):
+    rule_keys = {(s['start'], s['end'], s['label']) for s in rule_spans}
+    model_keep = decode(model_spans, threshold)
+    extra = sum(1 for s in model_keep if (s['start'], s['end'], s['label']) not in rule_keys)
+    _, rules_chars, _, _ = apply_union(text, rule_spans)
+    masked, union_chars, n_phone, n_pii = apply_union(text, list(rule_spans) + model_keep)
+    return masked, {
+        'phone': n_phone,
+        'pii': n_pii,
+        'rules_placeholder_chars': rules_chars,
+        'union_placeholder_chars': union_chars,
+        'model_extra_spans': extra,
+    }
+
+
+def _resolve(source, *, local_files_only):
+    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))
+
+
+def _load_rules_module(asset):
+    path = Path(asset) / 'scrub_pii.py'
+    if path.is_file():
+        import importlib.util
+        spec = importlib.util.spec_from_file_location('_nergal_scrub_pii', path)
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+        verify_rules(path)
+        return module
+    verify_rules()
+    return scrub_pii
+
+
+class Nergal:
+    def __init__(self, asset, device='cpu'):
+        import torch
+        from transformers import AutoModelForTokenClassification, AutoTokenizer
+        self.device = device
+        self._torch = torch
+        asset = Path(asset)
+        self._scrub = _load_rules_module(asset)
+        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')
+        tokenizer = AutoTokenizer.from_pretrained(
+            str(asset), local_files_only=True, use_fast=True, fix_mistral_regex=False,
+        )
+        if [tokenizer.convert_tokens_to_ids(t) for t in GAPS] != GAP_IDS:
+            raise ValueError('Packed NERGAL tokenizer is missing gap ids')
+        self.model = AutoModelForTokenClassification.from_pretrained(str(asset), local_files_only=True)
+        self.encoding = Encoding(tokenizer)
+        self.threshold = THRESHOLD
+        self.model.to(device).eval()
+
+    @classmethod
+    def from_pretrained(cls, source=HUB_ID, *, device=None, 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)
+
+    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():
+            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()}
+                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())
+
+    def rule_spans(self, text):
+        result = []
+        self._scrub.scrub_pii(text, spans=result)
+        if '[PII]' in text or '[Telefon]' in text:
+            return []
+        return sorted(({k: s[k] for k in ('start', 'end', 'label')} | {'score': 1.0} for s in result),
+                      key=lambda s: s['start'])
+
+    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)
+
+
+def main(argv=None):
+    import argparse
+    import sys
+    parser = argparse.ArgumentParser(description='NERGAL hybrid PII cleaner')
+    parser.add_argument('--repo', default=HUB_ID)
+    parser.add_argument('--device', default=None)
+    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)
+    text = sys.stdin.read()
+    masked, counts = nergal.scrub(text)
+    sys.stdout.write(masked)
+    print(json.dumps(counts), file=sys.stderr)
+
+
+if __name__ == '__main__':
+    main()