Spaces:
Sleeping
Sleeping
| """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 | |
| 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 | |