Instructions to use SlayerLab/NERGAL with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Transformers
How to use SlayerLab/NERGAL with Transformers:
# Use a pipeline as a high-level helper from transformers import pipeline pipe = pipeline("token-classification", model="SlayerLab/NERGAL")# pip install -U transformers accelerate # Load model directly from transformers import AutoTokenizer, AutoModelForTokenClassification tokenizer = AutoTokenizer.from_pretrained("SlayerLab/NERGAL") model = AutoModelForTokenClassification.from_pretrained("SlayerLab/NERGAL", device_map="auto") - Notebooks
- Google Colab
- Kaggle
Download test_nergal.py from SlayerLab/NERGAL: direct link, hf CLI and curl.
- Browser
- Download file 27.2 kB
-
https://huggingface.co/SlayerLab/NERGAL/resolve/main/test_nergal.py
- Command line
-
hf download hf://SlayerLab/NERGAL/test_nergal.py
-
curl -L -o test_nergal.py https://huggingface.co/SlayerLab/NERGAL/resolve/main/test_nergal.py
27.2 kB
| """Synthetic NERGAL tests. Invented strings only; no corpus text or real identifiers.""" | |
| import hashlib | |
| import json | |
| import re | |
| import shutil | |
| import tempfile | |
| import unittest | |
| from pathlib import Path | |
| HERE = Path(__file__).resolve().parent | |
| RULES_SHA = '4d0547238131908ca498a23607a7877512b2fb6022c484171334af1e604fc95a' | |
| # Procedural party-role nouns in every inflection the names layer leaves out, built here, not from ROLE_NOUNS | |
| PARTICIPLE_ENDINGS = ('y', 'ego', 'emu', 'ym', 'ych', 'ymi', 'a', 'ej', 'ą', 'e') | |
| ROLE_NOUN_FORMS = tuple(stem + e for stem in ('Odwołując', 'Przystępując', 'Zamawiając', 'Protestując', 'Skarżąc') | |
| for e in PARTICIPLE_ENDINGS) | |
| STEM_NAMES = ('Odwołującyński', 'Zamawiającek', 'Skarżącki') | |
| class NergalTests(unittest.TestCase): | |
| def test_card_and_rules_hash(self): | |
| 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, '2.0.4') | |
| self.assertEqual(card['version'], VERSION) | |
| self.assertEqual(card['eval']['union_fp'], 80) | |
| self.assertEqual(card['eval']['rules_fp'], 24) | |
| self.assertEqual((card['eval']['whole_entities'], card['eval']['gold_entities']), (303, 315)) # restated gold | |
| self.assertEqual(GAPS, card['gaps']) | |
| self.assertEqual(GAP_IDS, card['gap_ids']) | |
| self.assertEqual(THRESHOLD, card['threshold']) | |
| self.assertEqual(PINNED, RULES_SHA) | |
| digest = hashlib.sha256((HERE / 'scrub_pii.py').read_bytes()).hexdigest() | |
| self.assertEqual(digest, RULES_SHA) | |
| def test_real_tokenizer_preserves_batch_and_unit_alignment(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) | |
| words = ['A', '[PII_SPACE]', '1'] | |
| encoded, first = encoding.encode(words) | |
| self.assertIsInstance(encoded['input_ids'][0], list) | |
| 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 | |
| for marker in ('[PHONE]', '[Telefon]', '[PII]', '[PERSON]'): | |
| with self.subTest(marker=marker): | |
| text = f'Kontakt {marker}, NIP 1234567802.' # invented, checksum-valid | |
| [span] = rules(text) | |
| self.assertEqual(text[span['start']:span['end']], '1234567802') | |
| self.assertEqual(rules('a [PII] b [PHONE] c [PERSON] d [Telefon] e'), []) | |
| def test_phones_are_tagged_phone(self): | |
| from nergal import apply_union, rules | |
| text = 'Biuro: (22) 123 45 67.' | |
| masked = apply_union(text, rules(text))[0] | |
| self.assertEqual(masked, 'Biuro: [PHONE].') | |
| def test_new_and_legacy_phone_tags_are_the_same_boundary(self): | |
| from nergal import rules | |
| def found(text): # values, not offsets: the two tags differ in length | |
| return [(text[s['start']:s['end']], s['label']) for s in rules(text)] | |
| for text in ('Telefon: {} lub 601234567', 'Kontakt: {}, 601234567', # plain 9 digits: cue-gated | |
| 'tel. {}\nwew. 123 Jan Nowak\nwew. 456 sekretariat', | |
| 'Kontakt {}: e-mail biuro@example.pl, 601 234 567'): # invented | |
| with self.subTest(text=text): | |
| self.assertEqual(found(text.format('[Telefon]')), found(text.format('[PHONE]'))) | |
| def test_a_name_placeholder_ends_an_other_number_label_but_never_a_phone_cue(self): | |
| from nergal import rules | |
| for text, masked in (('Kod [PERSON] zadzwoń 601 200 300', ['601 200 300']), | |
| ('Kod Jan zadzwoń 601 200 300', []), # a raw name is text to the rules | |
| ('[PERSON] NIP: 601 234 567', []), # a label after the name still counts | |
| ('Telefon do [PERSON]: 601234567', ['601234567']), # plain 9 digits: cue-gated | |
| ('Kontakt: [PERSON], 601234567', ['601234567']), | |
| ('Informacje u [PERSON] pod numerem 601234567', ['601234567']), | |
| ('tel. [PERSON] 601234567', ['601234567'])): # invented | |
| with self.subTest(text=text): | |
| self.assertEqual([text[s['start']:s['end']] for s in rules(text) if s['label'] == 'phone'], masked) | |
| def test_a_placeholder_between_a_label_and_its_value_or_inside_a_value(self): | |
| from nergal import apply_union, rules | |
| from scrub_pii import scrub_pii | |
| for text, masked in ( # invented; 5260001246 and 123456785 are checksum-valid | |
| ('NIP [PII] 526-000-12-46', 'NIP [PII] [PII]'), # the label reaches past one placeholder | |
| ('NIP: [PII], 526-000-12-46.', 'NIP: [PII], [PII].'), | |
| ('REGON [PII] 123456785', 'REGON [PII] [PII]'), | |
| ('NIP [PII]\n5260001246', 'NIP [PII]\n[PII]'), | |
| ('NIP [PII] 526-000-12-47', None), # still checksum-gated | |
| ('NIP [PII] [PII] 5260001246', None), # one placeholder only | |
| ('NIP 526-000-12-46, 5260001246', 'NIP [PII], [PII]'), # this scrub's own placeholder counts too | |
| ('NIP 526-000-12-46 i 5260001246', 'NIP [PII] i 5260001246'), | |
| ('PESEL 8503[PII]1212343', 'PESEL [PII][PII][PII]'), # labelled fragments are masked | |
| ('NIP 526-000-[PII]-46', 'NIP [PII]-[PII]-[PII]'), | |
| ('PESEL: 85031[PII]', 'PESEL: [PII][PII]'), | |
| ('PESEL 1234 [PII]', None), # under 5 digits in all | |
| ('tel. 601 [PII] 200 300', 'tel. [PHONE] [PII] [PHONE]'), # a phone cue: fragments are [PHONE] | |
| ('tel. 601 200 [PHONE]', 'tel. [PHONE] [PHONE]'), | |
| ('tel. +48 601 [PII] 300', 'tel. [PHONE] [PII] [PHONE]'), | |
| ('Kontakt: 22 123 [PHONE] 67', 'Kontakt: [PHONE] [PHONE] [PHONE]'), | |
| ('8503[PII]1212343', None), # no label, no cue: known gap | |
| ('601 [PII] 200 300', None), ('601 [PHONE] 200 300', None), | |
| ('tel. 601 200 [PHONE]abc', None), # glued to a word | |
| ('tel. 00-950 [PHONE]', None), ('tel. 2023-10-16 [PHONE]', None), # postcode, date | |
| ('tel. 12 345 [PII] zł', None), # amount | |
| ('Strona [PII] 12 z 40', None), ('[PII] 2019 r.', None), ('art. [PII] ust. 2', None), | |
| ('tel. [PHONE], godz. 8–16', None), ('[PHONE] 123 osób', None), | |
| ('Kontakt: [PHONE], 2 000 000 zł', None), ('NIP [PII], rok 2023', None), | |
| ('tel. [PHONE], pok. 101', None), ('tel. [PHONE], kod 00-950', None), ('[PII] 12 345 678 zł', None), | |
| ('tel. [PHONE] w godz. 10–18, sala 204', None)): | |
| with self.subTest(text=text): | |
| spans = rules(text) | |
| self.assertEqual(apply_union(text, spans)[0], masked or text) | |
| self.assertTrue(all('[' not in text[s['start']:s['end']] for s in spans)) # placeholders stay as they are | |
| self.assertEqual(scrub_pii('PESEL 8503[PII]1212343')[1]['pesel'], 1) | |
| self.assertEqual(scrub_pii('tel. 601 [PII] 200 300')[1]['phone'], 1) | |
| def test_grouped_national_phones_mask_without_a_cue(self): | |
| from nergal import rules | |
| for text, masked in (('Sklep Ala, 601 234 567, czynne 9-17', ['601 234 567']), | |
| ('Biuro: (22) 123 45 67.', ['(22) 123 45 67']), | |
| ('Zapraszamy: +48 601 234 567.', ['+48 601 234 567']), | |
| ('Zapraszamy: 601234567.', []), # plain 9 digits stay cue-gated | |
| ('Budżet wyniósł 601 234 567 zł.', []), # amount | |
| ('Wartość 601 234 567,89 w tabeli.', []), # decimal figure | |
| ('Kwota 500 000 000 osób.', []), # round count | |
| ('NIP: 601 234 567', [])): # other-number label | |
| with self.subTest(text=text): | |
| self.assertEqual([text[s['start']:s['end']] for s in rules(text) if s['label'] == 'phone'], masked) | |
| def test_email_ends_at_glued_text_and_mention_lists_are_not_addresses(self): | |
| from nergal import rules | |
| for text, masked in (('kontakt@firma.plKontakt', ['kontakt@firma.pl']), # capital glued onto the TLD | |
| ('jan@firma.plwww.firma.pl', ['jan@firma.pl']), # URL host glued onto the TLD | |
| ('jan@firma.plkontakt', ['jan@firma.plkontakt']), # all-lowercase glue: known limit | |
| ('BIURO@FIRMA.PL', ['BIURO@FIRMA.PL']), | |
| ('kontakt@jan7@wp.pl', ['jan7@wp.pl']), # word glued on with '@' | |
| ('Dzięki @kasia @firma.pl @tomek', []), # list of mentions | |
| ('Obserwuj @jan@firma.social', ['jan@firma.social'])): # handle keeps the mask | |
| with self.subTest(text=text): | |
| self.assertEqual([text[s['start']:s['end']] for s in rules(text)], masked) | |
| def test_union_keeps_regex_and_adds_model_spans(self): | |
| from nergal import apply_union, scrub_spans | |
| text = 'Ring 000000000 then extra.' | |
| rules = [{'start': 5, 'end': 14, 'label': 'phone', 'score': 1.0}] | |
| model = [ | |
| {'start': 5, 'end': 14, 'label': 'phone', 'score': 0.99}, | |
| {'start': 20, 'end': 25, 'label': 'pii', 'score': 0.97}, | |
| ] | |
| masked, counts = scrub_spans(text, rules, model, threshold=0.95) | |
| self.assertIn('[PHONE]', masked) | |
| 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) | |
| 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_a_rule_label_wins_on_its_characters(self): | |
| from nergal import char_labels | |
| spans = [{'start': 0, 'end': 4, 'label': 'phone'}, {'start': 2, 'end': 6, 'label': 'pii'}, | |
| {'start': 5, 'end': 8, 'label': 'phone'}] | |
| self.assertEqual(char_labels('x' * 8, spans), ['phone'] * 4 + ['pii'] + ['phone'] * 3) | |
| self.assertEqual(char_labels('x' * 8, spans, firm=spans[1:2]), ['phone'] * 2 + ['pii'] * 4 + ['phone'] * 2) | |
| self.assertEqual(char_labels('x' * 8, spans, firm=spans), ['phone'] * 4 + ['pii'] + ['phone'] * 3) # firm vs firm: rank | |
| self.assertEqual(char_labels('xy', []), [None, None]) | |
| def test_labelled_identifiers_keep_pii_under_a_model_phone(self): | |
| from nergal import rules, scrub_spans | |
| account = '28 1020 5558 1111 1119 0100 0000' # invented, checksum-valid | |
| for value, label in (('44051401359', 'PESEL'), ('44051401359', 'PESEL:'), ('123456785', 'REGON'), | |
| ('12345678512347', 'REGON'), ('1234567802', 'NIP'), ('123-456-78-02', 'NIP'), | |
| ('ABA300000', 'dowód osobisty'), ('AB1234567', 'paszport nr'), (account, 'konto'), | |
| (account.replace(' ', ''), 'nr rachunku')): | |
| for lead in ('Kontakt: tel. 601 222 333, ', 'Dane: tel. 601 222 333, ', 'tel. 601 222 333, ', 'Kontakt: ', ''): | |
| text = f'{lead}{label} {value}.' | |
| at = text.index(value, len(lead)) | |
| phone = [dict(start=at, end=at + len(value), label='phone', score=0.99)] | |
| with self.subTest(text=text): | |
| found = rules(text) | |
| self.assertFalse(any(s.get('bare') for s in found)) | |
| masked = scrub_spans(text, found, phone)[0] | |
| self.assertTrue(masked.endswith(f'{label} [PII].'), masked) | |
| self.assertEqual(masked.count('[PHONE]'), lead.count('601')) # the contact phone keeps [PHONE] | |
| def test_an_unlabelled_pesel_yields_to_a_model_phone(self): | |
| from nergal import rules, scrub_spans | |
| for text in ('Kontakt: 44051401359', 'tel. 44051401359', 'PESEL? Nie, tel. 44051401359.'): | |
| at = next(i for i, c in enumerate(text) if c.isdigit()) | |
| phone = [dict(start=at, end=at + 11, label='phone', score=0.99)] | |
| with self.subTest(text=text): | |
| found = rules(text) | |
| self.assertTrue(found and all(s.get('bare') and s['label'] == 'pii' for s in found)) # rules alone: [PII] | |
| self.assertEqual(scrub_spans(text, found, [])[0], text[:at] + '[PII]' + text[at + 11:]) | |
| self.assertEqual(scrub_spans(text, found, phone)[0], text[:at] + '[PHONE]' + text[at + 11:]) | |
| def test_a_rule_phone_keeps_phone_under_a_model_pii(self): | |
| from nergal import rules, scrub_spans | |
| text = 'Kontakt: tel. 601 222 333, kom. 602 333 444.' | |
| model = [dict(start=text.index('602'), end=text.index('602') + 11, label='pii', score=0.99)] | |
| self.assertEqual(scrub_spans(text, rules(text), model)[0], 'Kontakt: tel. [PHONE], kom. [PHONE].') | |
| 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_person_spans_leave_procedural_role_nouns_out(self): | |
| from nergal import person_spans | |
| def persons(text): # every word tagged, as the names model tags a role noun next to a name | |
| return [text[s['start']:s['end']] for s in person_spans(text, [m.span() for m in re.finditer(r'\S+', text)])] | |
| for noun in ROLE_NOUN_FORMS: | |
| self.assertEqual(persons(noun), [], noun) | |
| self.assertEqual(persons(noun.upper()), [], noun) | |
| self.assertEqual(persons('Odwołujący Jan Nowak'), ['Jan Nowak']) # the name stays | |
| self.assertEqual(persons('Jan Nowak Przystępujący Anna Kowal'), ['Jan Nowak', 'Anna Kowal']) # the noun splits | |
| text = 'od Przystępującego, na rzecz' | |
| comma = text.index(',') | |
| self.assertEqual(person_spans(text, [(comma, comma + 1)]), []) # a tagged comma grows to the noun | |
| self.assertEqual(persons('(Zamawiającemu)'), []) | |
| for name in STEM_NAMES: # invented surnames that start like a role noun | |
| self.assertEqual(persons(name), [name]) | |
| 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( | |
| start=text.index(part), end=text.index(part) + len(part), label=label, score=score) | |
| for text, part, kept in (('tel. 112', '112', []), # emergency number | |
| ('tel. 51 23 45', '51 23 45', []), # under 7 digits | |
| ('tel. 601 234 567/602 345 678', '601 234 567/602 345 678', | |
| ['601 234 567', '602 345 678']), # one span per number | |
| ('tel. 601 234 567, fax, 602 345 678', '601 234 567, fax, 602 345 678', | |
| ['601 234 567', '602 345 678']), # a word between parts | |
| ('tel. 22 123 45 67 wew. 101', '22 123 45 67 wew. 101', | |
| ['22 123 45 67 wew. 101']), # extension stays inside | |
| ('Jan Kowalski, 112', 'Jan Kowalski', ['Jan Kowalski'])): # other labels unchanged | |
| label = 'pii' if part[0].isalpha() else 'phone' | |
| with self.subTest(text=text): | |
| keep = model_keep(text, [span(text, part, label)]) | |
| self.assertEqual([text[s['start']:s['end']] for s in keep], kept) | |
| self.assertTrue(all(s['label'] == label and s['score'] == 0.99 for s in keep)) | |
| text = 'tel. 112' | |
| 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 test_published_weights_label_contact_lines(self): | |
| from nergal import Nergal | |
| nergal = Nergal(HERE, device='cpu') | |
| account = '28 1020 5558 1111 1119 0100 0000' # invented, checksum-valid | |
| cases = { | |
| 'Kontakt: tel. 601 222 333, PESEL 44051401359.': 'Kontakt: tel. [PHONE], PESEL [PII].', | |
| 'Kontakt: jan.kowalski@example.com, PESEL 44051401359.': 'Kontakt: [PII], PESEL [PII].', | |
| 'Kontakt: tel. 601 222 333, REGON 123456785.': 'Kontakt: tel. [PHONE], REGON [PII].', | |
| 'Kontakt: tel. 601 222 333, REGON 12345678512347.': 'Kontakt: tel. [PHONE], REGON [PII].', | |
| 'Kontakt: PESEL 44051401359.': 'Kontakt: PESEL [PII].', | |
| 'Dane: tel. 601 222 333, PESEL 44051401359.': 'Dane: tel. [PHONE], PESEL [PII].', | |
| 'tel. 601 222 333, PESEL 44051401359.': 'tel. [PHONE], PESEL [PII].', | |
| 'Telefon: 601 222 333, PESEL: 44051401359': 'Telefon: [PHONE], PESEL: [PII]', | |
| 'PESEL 44051401359.': 'PESEL [PII].', | |
| 'Kontakt: tel. 601 222 333, NIP 1234567802.': 'Kontakt: tel. [PHONE], NIP [PII].', | |
| 'Kontakt: tel. 601 222 333, NIP 123-456-78-02.': 'Kontakt: tel. [PHONE], NIP [PII].', | |
| 'Kontakt: tel. 601 222 333, dowód osobisty ABA300000.': 'Kontakt: tel. [PHONE], dowód osobisty [PII].', | |
| 'Kontakt: tel. 601 222 333, paszport nr AB1234567.': 'Kontakt: tel. [PHONE], paszport nr [PII].', | |
| f'Kontakt: tel. 601 222 333, konto {account}.': 'Kontakt: tel. [PHONE], konto [PII].', | |
| 'Kontakt: tel. 601 222 333, kom. 602 333 444.': 'Kontakt: tel. [PHONE], kom. [PHONE].', | |
| 'Kontakt: +48 22 123 45 67, e-mail biuro@example.pl.': 'Kontakt: [PHONE], e-mail [PII].', | |
| 'Dane: tel. (22) 123 45 67, fax 22 123 45 68.': 'Dane: tel. [PHONE], fax [PHONE].', | |
| 'Kontakt: 44051401359': 'Kontakt: [PHONE]', # an unlabelled PESEL-valid number the model reads as a phone | |
| 'Zadzwoń: 48601222333.': 'Zadzwoń: [PHONE].', | |
| } | |
| self.assertEqual({t: m for t, (m, _) in zip(cases, nergal.scrub_many(list(cases)))}, cases) | |
| def test_published_weights_mask_values_split_by_a_placeholder(self): | |
| from nergal import Nergal | |
| nergal = Nergal(HERE, device='cpu') | |
| cases = { # invented; the model adds nothing to the rules here, and the gaps stay unmasked | |
| 'NIP [PII] 526-000-12-46': 'NIP [PII] [PII]', 'NIP: [PII], 526-000-12-46.': 'NIP: [PII], [PII].', | |
| 'REGON [PII] 123456785': 'REGON [PII] [PII]', 'PESEL 8503[PII]1212343': 'PESEL [PII][PII][PII]', | |
| 'NIP 526-000-[PII]-46': 'NIP [PII]-[PII]-[PII]', 'tel. 601 [PII] 200 300': 'tel. [PHONE] [PII] [PHONE]', | |
| 'tel. 601 200 [PHONE]': 'tel. [PHONE] [PHONE]', 'tel. [Telefon], 601 200 300': 'tel. [Telefon], [PHONE]', | |
| **{t: t for t in ('8503[PII]1212343', '601 [PII] 200 300', '601 [PHONE] 200 300', 'Strona [PII] 12 z 40', | |
| '[PII] 2019 r.', 'art. [PII] ust. 2', 'tel. [PHONE], godz. 8–16', '[PHONE] 123 osób', | |
| 'Kontakt: [PHONE], 2 000 000 zł', 'NIP [PII], rok 2023', 'tel. [PHONE], pok. 101', | |
| 'tel. [PHONE], kod 00-950', '[PII] 12 345 678 zł', 'tel. [PHONE] w godz. 10–18, sala 204')}, | |
| } | |
| self.assertEqual({t: m for t, (m, _) in zip(cases, nergal.scrub_many(list(cases)))}, cases) | |
| 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_leave_role_nouns_out_and_keep_the_names(self): | |
| names = self._names() | |
| texts = ['W toku rozprawy Przystępujący podtrzymał swoje stanowisko. Odwołujący wniósł o uwzględnienie odwołania, ' | |
| 'a Zamawiający o jego oddalenie.', | |
| 'Izba oddaliła odwołanie. Odwołujący nie wykazał, że Przystępującemu należało odrzucić ofertę.', | |
| 'Odwołujący Jan Wrzos wniósł odwołanie. Przystępujący Tomasz Kmieć zgłosił przystąpienie po stronie ' | |
| 'Zamawiającego.', | |
| 'Protestująca wniosła protest. Zamawiający oddalił protest Protestującej, a Protestująca wniosła odwołanie.', | |
| 'Skarżąca wniosła skargę na decyzję. Organ wniósł o oddalenie skargi Skarżącej.'] | |
| found = [[t[s['start']:s['end']] for s in spans] for t, spans in zip(texts, names.spans_many(texts))] | |
| # 2.0.3: 'Przystępujący', 'Przystępującemu', 'Protestującej', 'Protestująca', 'Skarżącej' | |
| self.assertEqual(found, [[], [], ['Jan Wrzos', 'Tomasz Kmieć'], [], []]) | |
| 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() | |