File size: 2,638 Bytes
4039e77
 
 
 
 
 
 
15eb2b4
4039e77
 
 
 
7ce967b
4039e77
 
15eb2b4
7ce967b
 
 
4039e77
 
 
 
 
 
 
510b295
 
 
 
 
 
 
 
 
 
 
4039e77
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
"""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 = '3016ae5bd403ff997458f9dd74bad8c6ed1388eb83dadc1b31cdb182f9ed607f'


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.0.2')
        self.assertEqual(card['version'], VERSION)
        self.assertEqual(card['eval']['union_fp'], 123)
        self.assertEqual(card['eval']['rules_fp'], 98)
        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_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()