"""Tests for subword–label alignment — the highest-risk code in the project. These tests use hand-constructed cases with known expected outputs. They must all pass before any training is attempted. Cases covered: - Simple single-word entities (no subwords) - Multi-subword tokens (e.g. numbers, company names) - Punctuation and adjacent tokens - Financial numbers like "$1.2B" (often split into multiple subwords) - B- label NOT promoted to I- at subword continuations - -100 (IGNORE_INDEX) on all continuation subwords - Truncation: labels don't leak across boundary - Empty / all-O sequence """ import sys from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src")) import pytest from transformers import AutoTokenizer from finner.labels import LABEL2ID, IGNORE_INDEX MODEL = "bert-base-uncased" # fast to load; same subword logic as DeBERTa @pytest.fixture(scope="module") def tokenizer(): return AutoTokenizer.from_pretrained(MODEL, use_fast=True) def align(tokenizer, tokens, ner_tags): from finner.data.align import align_labels_to_subwords return align_labels_to_subwords(tokenizer, tokens, ner_tags) class TestBasicAlignment: def test_single_word_entity(self, tokenizer): tokens = ["Apple", "reported", "revenue"] tags = ["B-ORG", "O", "O"] result = align(tokenizer, tokens, tags) labels = result["labels"] # Must have same length as input_ids assert len(labels) == len(result["input_ids"]) # Only non-IGNORE_INDEX labels should map to known label ids real_labels = [l for l in labels if l != IGNORE_INDEX] assert LABEL2ID["B-ORG"] in real_labels assert LABEL2ID["O"] in real_labels def test_all_o_sequence(self, tokenizer): tokens = ["The", "market", "opened", "higher"] tags = ["O", "O", "O", "O"] result = align(tokenizer, tokens, tags) real_labels = [l for l in result["labels"] if l != IGNORE_INDEX] assert all(l == LABEL2ID["O"] for l in real_labels) def test_label_count_matches_input_ids(self, tokenizer): tokens = ["Goldman", "Sachs", "paid", "$500M", "to", "the", "SEC"] tags = ["B-ORG", "I-ORG", "O", "B-MONEY","O", "O", "B-REG_BODY"] result = align(tokenizer, tokens, tags) assert len(result["labels"]) == len(result["input_ids"]) def test_attention_mask_length(self, tokenizer): tokens = ["Revenue", "grew", "12.5%"] tags = ["B-METRIC", "O", "B-PERCENT"] result = align(tokenizer, tokens, tags) assert len(result["attention_mask"]) == len(result["input_ids"]) class TestSubwordAlignment: def test_multisubword_entity_first_subword_only(self, tokenizer): """For a word split into N subwords, only the first gets the real label.""" # "$1.2B" is very likely split by BERT's tokenizer tokens = ["$1.2B"] tags = ["B-MONEY"] result = align(tokenizer, tokens, tags) # Gather non-special-token labels word_ids = tokenizer(tokens, is_split_into_words=True).word_ids() labels = result["labels"] first_seen = False for word_id, label in zip(word_ids, labels): if word_id == 0: if not first_seen: assert label == LABEL2ID["B-MONEY"], "First subword must get B-MONEY" first_seen = True else: assert label == IGNORE_INDEX, "Continuation subwords must be IGNORE_INDEX" def test_b_label_not_promoted_to_i(self, tokenizer): """B-X on word boundary must remain B-X, not become I-X on any subword.""" tokens = ["Apple", "Inc", "reported"] tags = ["B-ORG", "I-ORG", "O"] result = align(tokenizer, tokens, tags) word_ids = tokenizer(tokens, is_split_into_words=True).word_ids() labels = result["labels"] # First subword of "Apple" (word_id=0) must be B-ORG first_apple = next(i for i, wid in enumerate(word_ids) if wid == 0) assert labels[first_apple] == LABEL2ID["B-ORG"] # First subword of "Inc" (word_id=1) must be I-ORG first_inc = next(i for i, wid in enumerate(word_ids) if wid == 1) assert labels[first_inc] == LABEL2ID["I-ORG"] def test_continuation_subwords_are_ignored(self, tokenizer): """All non-first subwords must be IGNORE_INDEX regardless of entity type.""" tokens = ["Massachusetts"] # likely multi-subword tags = ["B-GPE"] result = align(tokenizer, tokens, tags) word_ids = tokenizer(tokens, is_split_into_words=True).word_ids() labels = result["labels"] first_seen = False for word_id, label in zip(word_ids, labels): if word_id is None: continue if not first_seen: assert label != IGNORE_INDEX, "First subword must have a real label" first_seen = True else: assert label == IGNORE_INDEX, "All continuation subwords must be IGNORE_INDEX" class TestSpecialTokens: def test_cls_sep_are_ignore(self, tokenizer): tokens = ["revenue"] tags = ["B-METRIC"] result = align(tokenizer, tokens, tags) word_ids = tokenizer(tokens, is_split_into_words=True).word_ids() labels = result["labels"] for word_id, label in zip(word_ids, labels): if word_id is None: assert label == IGNORE_INDEX, "Special tokens [CLS]/[SEP] must be IGNORE_INDEX" class TestTruncation: def test_truncation_does_not_leak_labels(self, tokenizer): """With a very short max_length, truncated tokens must not appear as labels.""" from finner.data.align import align_labels_to_subwords tokens = ["Apple", "reported", "net", "income", "of", "$1.2B", "in", "Q3"] tags = ["B-ORG", "O", "O", "B-METRIC","O", "B-MONEY","O", "B-DATE"] result = align_labels_to_subwords(tokenizer, tokens, tags, max_length=8) # Labels must exactly match input_ids in length assert len(result["labels"]) == len(result["input_ids"]) assert len(result["input_ids"]) <= 8 class TestFinancialEdgeCases: def test_ticker_symbol(self, tokenizer): tokens = ["$AAPL", "rose", "3%"] tags = ["B-TICKER", "O", "B-PERCENT"] result = align(tokenizer, tokens, tags) assert len(result["labels"]) == len(result["input_ids"]) def test_monetary_with_suffix(self, tokenizer): tokens = ["$", "500", "million"] tags = ["B-MONEY", "I-MONEY", "I-MONEY"] result = align(tokenizer, tokens, tags) real = [l for l in result["labels"] if l != IGNORE_INDEX] assert LABEL2ID["B-MONEY"] in real