2.0.0: ranked union, person spans

PHONE > PII > PERSON on overlap; person token ranges expand to whole words
and merge into spans; scrub_spans counts person. Names off: output unchanged.

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

nergal.py

@@ -14,7 +14,7 @@ from functools import lru_cache
 from pathlib import Path
 
 import scrub_pii
-from scrub_pii import PHONE_TAG, PII_TAG
+from scrub_pii import LEGACY_PHONE_TAG, PHONE_TAG, PII_TAG
 
 HUB_ID = 'SlayerLab/NERGAL'
 VERSION = '1.2.0'
@@ -25,6 +25,13 @@ LABELS = ['phone', 'pii']
 THRESHOLD = 0.95
 RULES_SHA = 'd18662434b3bab69d74e7f63f76a646122cd9d7420de1d70914aaeb903ff2357'
 DTYPES = ('float32', 'float16')
+PERSON_TAG = '[PERSON]'
+MARKERS = (PHONE_TAG, PII_TAG, PERSON_TAG, LEGACY_PHONE_TAG)
+TAGS = {'person': PERSON_TAG, 'pii': PII_TAG, 'phone': PHONE_TAG}  # ascending precedence
+RANK = {label: i for i, label in enumerate(TAGS)}
+PERSON_LABELS = ('PERSON', 'PERSON_F', 'PERSON_L')
+_WORD = re.compile(r"[\w'’-]")
+_PLACEHOLDER = re.compile('|'.join(map(re.escape, MARKERS)))
 
 
 def sha(path):
@@ -180,16 +187,43 @@ def rules(text):
     return spans_from(scrub_pii, text)
 
 
+def expand(text, start, end):
+    """Grow a token range to the whole word: letters, digits, apostrophes and hyphens."""
+    while start > 0 and _WORD.match(text, start - 1):
+        start -= 1
+    while end < len(text) and _WORD.match(text, end):
+        end += 1
+    return start, end
+
+
+def person_spans(text, ranges):
+    """Person spans from FastPDN person-token ranges: whole words, joined across spaces, tabs and no-break spaces
+    (not punctuation or line breaks); words touching a placeholder are skipped."""
+    blocked = [m.span() for m in _PLACEHOLDER.finditer(text)]
+    merged = []
+    for start, end in sorted(expand(text, a, b) for a, b in ranges):
+        if any(start < b and a < end for a, b in blocked):
+            continue
+        if merged and not text[merged[-1][1]:start].strip(' \t\u00a0'):
+            merged[-1][1] = max(merged[-1][1], end)
+        else:
+            merged.append([start, end])
+    return [{'start': a, 'end': b, 'label': 'person', 'score': 1.0} for a, b in merged]
+
+
 def apply_union(text, spans):
+    """Each character takes its highest-ranked label (phone > pii > person); each run becomes one placeholder."""
     labels = [None] * len(text)
     for span in spans:
         start, end, label = span['start'], span['end'], span['label']
+        if label not in RANK:
+            raise ValueError(f'Unknown label {label!r}')
         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':
+            if labels[i] is None or RANK[label] > RANK[labels[i]]:
                 labels[i] = label
-    out, chars, n_phone, n_pii, i = [], 0, 0, 0, 0
+    out, chars, counts, i = [], 0, dict.fromkeys(TAGS, 0), 0
     while i < len(text):
         lab = labels[i]
         if lab is None:
@@ -199,15 +233,11 @@ def apply_union(text, spans):
         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
+        out.append(TAGS[lab])
+        chars += len(TAGS[lab])
+        counts[lab] += 1
         i = j
-    return ''.join(out), chars, n_phone, n_pii
+    return ''.join(out), chars, counts
 
 
 def model_keep(text, model_spans, *, threshold=THRESHOLD, module=scrub_pii):
@@ -218,19 +248,14 @@ def model_keep(text, model_spans, *, threshold=THRESHOLD, module=scrub_pii):
                          else [(s['start'], s['end'])])]
 
 
-def scrub_spans(text, rule_spans, model_spans, *, threshold=THRESHOLD, module=scrub_pii):
+def scrub_spans(text, rule_spans, model_spans, *, threshold=THRESHOLD, module=scrub_pii, name_spans=()):
     rule_keys = {(s['start'], s['end'], s['label']) for s in rule_spans}
     keep = model_keep(text, model_spans, threshold=threshold, module=module)
     extra = sum(1 for s in 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) + keep)
-    return masked, {
-        'phone': n_phone,
-        'pii': n_pii,
-        'rules_placeholder_chars': rules_chars,
-        'union_placeholder_chars': union_chars,
-        'model_extra_spans': extra,
-    }
+    _, rules_chars, _ = apply_union(text, rule_spans)
+    masked, union_chars, counts = apply_union(text, [*rule_spans, *keep, *name_spans])
+    return masked, {**counts, 'rules_placeholder_chars': rules_chars, 'union_placeholder_chars': union_chars,
+                    'model_extra_spans': extra}
 
 
 def _resolve(source, *, local_files_only):

test_nergal.py

@@ -119,11 +119,37 @@ class NergalTests(unittest.TestCase):
         self.assertIn('[PII]', masked)
         self.assertGreater(counts['union_placeholder_chars'], counts['rules_placeholder_chars'])
         self.assertEqual(counts['model_extra_spans'], 1)
-        _, rules_chars, _, _ = apply_union(text, rules)
+        _, rules_chars, _ = apply_union(text, rules)
         self.assertEqual(counts['rules_placeholder_chars'], rules_chars)
+        self.assertEqual(counts['person'], 0)
         self.assertNotIn('000000000', masked)
         self.assertNotIn('extra', masked)
 
+    def test_union_ranks_phone_over_pii_over_person(self):
+        from nergal import apply_union
+        text = 'Jan Kowal 601 234 567'
+        spans = [{'start': 0, 'end': 21, 'label': 'person'}, {'start': 4, 'end': 9, 'label': 'pii'},
+                 {'start': 10, 'end': 21, 'label': 'phone'}]
+        masked, chars, counts = apply_union(text, spans)
+        self.assertEqual(masked, '[PERSON][PII][PERSON][PHONE]')  # 0–4 and the space at 9 stay person
+        self.assertEqual(chars, len(masked))
+        self.assertEqual(counts, {'person': 2, 'pii': 1, 'phone': 1})
+        with self.assertRaises(ValueError):
+            apply_union(text, [{'start': 0, 'end': 1, 'label': 'org'}])
+
+    def test_person_spans_expand_merge_and_skip_markers(self):
+        from nergal import person_spans
+        text = 'Pani Nowakowskiej-Kowal, O’Brien. [PERSON] i [PHONE]'
+        a = text.index('Nowak'); b = text.index('O’B')
+        spans = person_spans(text, [(a + 2, a + 6), (b, b + 2), (text.index('[PERSON]') + 1, text.index('[PERSON]') + 3)])
+        self.assertEqual([text[s['start']:s['end']] for s in spans], ['Nowakowskiej-Kowal', 'O’Brien'])
+        self.assertEqual(person_spans('Jan Nowak', [(0, 3), (4, 9)]), [{'start': 0, 'end': 9, 'label': 'person', 'score': 1.0}])
+        self.assertEqual(person_spans('Jan Nowak', [(0, 3), (4, 9)])[0]['end'], 9)  # no-break space joins
+        self.assertEqual(person_spans('Jan, Nowak', [(0, 3), (5, 10)])[1]['start'], 5)  # punctuation keeps them apart
+        self.assertEqual(len(person_spans('Jan\nNowak', [(0, 3), (4, 9)])), 2)  # a line break keeps them apart
+        self.assertEqual(person_spans('Nowak', [(0, 2), (1, 5)]), [{'start': 0, 'end': 5, 'label': 'person', 'score': 1.0}])
+        self.assertEqual(person_spans('x', []), [])
+
     def test_model_phone_spans_follow_the_phone_policy(self):
         from nergal import model_keep, scrub_spans
         span = lambda text, part, label='phone', score=0.99: dict(