test_nergal.py

"""Synthetic NERGAL tests. Invented strings only; no corpus text or real identifiers."""
import hashlib
import json
import unittest
from pathlib import Path

HERE = Path(__file__).resolve().parent
RULES_SHA = '08faef844c850bcd438c904d0b3f898df47c8bde8dd827d36ebffd39dc1594fb'


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, '1.1.2')
        self.assertEqual(card['version'], VERSION)
        self.assertEqual(card['eval']['union_fp'], 80)
        self.assertEqual(card['eval']['rules_fp'], 24)
        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
        text = 'Kontakt [Telefon], NIP 1234567802.'  # invented, checksum-valid
        [span] = rules(text)
        self.assertEqual(text[span['start']:span['end']], '1234567802')
        self.assertEqual(rules('a [PII] b [Telefon] c'), [])

    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('[Telefon]', 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.assertNotIn('000000000', masked)
        self.assertNotIn('extra', masked)


if __name__ == '__main__':
    unittest.main()