Spaces:
Sleeping
Sleeping
| """Tests for the inference pipeline (predict.py). | |
| These tests run without a trained checkpoint by using a randomly-initialized | |
| model, which is sufficient to verify the pipeline structure and output schema. | |
| """ | |
| import sys | |
| from pathlib import Path | |
| sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src")) | |
| import pytest | |
| import torch | |
| from unittest.mock import patch, MagicMock | |
| from finner.labels import ENTITY_TYPES, LABEL2ID | |
| class TestPredictOutputSchema: | |
| """Verify predict() returns well-formed output with correct keys and types.""" | |
| def mock_model(self, tmp_path): | |
| """Patch _load_model to return a tiny random model instead of a real checkpoint.""" | |
| from transformers import AutoModelForTokenClassification, AutoTokenizer | |
| from finner.labels import NUM_LABELS, ID2LABEL, LABEL2ID | |
| model_name = "bert-base-uncased" # already cached from training | |
| tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True) | |
| model = AutoModelForTokenClassification.from_pretrained( | |
| model_name, | |
| num_labels=NUM_LABELS, | |
| id2label=ID2LABEL, | |
| label2id=LABEL2ID, | |
| ignore_mismatched_sizes=True, | |
| ) | |
| model.eval() | |
| import finner.infer.predict as predict_module | |
| predict_module._model = model | |
| predict_module._tokenizer = tokenizer | |
| yield | |
| predict_module._model = None | |
| predict_module._tokenizer = None | |
| def test_output_keys_present(self): | |
| from finner.infer.predict import predict | |
| result = predict("Apple reported revenue of $1.2B.") | |
| assert "tokens" in result | |
| assert "token_labels" in result | |
| assert "token_confidences" in result | |
| assert "entities" in result | |
| def test_tokens_are_strings(self): | |
| from finner.infer.predict import predict | |
| result = predict("Goldman Sachs paid $500M to the SEC.") | |
| assert all(isinstance(t, str) for t in result["tokens"]) | |
| def test_label_count_matches_token_count(self): | |
| from finner.infer.predict import predict | |
| result = predict("Revenue grew 12.5% year over year.") | |
| assert len(result["token_labels"]) == len(result["tokens"]) | |
| assert len(result["token_confidences"]) == len(result["tokens"]) | |
| def test_all_labels_are_valid(self): | |
| from finner.infer.predict import predict | |
| result = predict("The Fed raised rates by 25 bps.") | |
| valid = set(LABEL2ID.keys()) | |
| for lbl in result["token_labels"]: | |
| assert lbl in valid, f"Invalid label: {lbl}" | |
| def test_confidences_in_unit_interval(self): | |
| from finner.infer.predict import predict | |
| result = predict("Microsoft acquired Activision for $68.7B.") | |
| for conf in result["token_confidences"]: | |
| assert 0.0 <= conf <= 1.0 | |
| def test_entity_schema(self): | |
| from finner.infer.predict import predict | |
| result = predict("Apple Inc. reported $1.2B EBITDA.") | |
| for ent in result["entities"]: | |
| assert "text" in ent | |
| assert "label" in ent | |
| assert "start_token" in ent | |
| assert "end_token" in ent | |
| assert "confidence" in ent | |
| assert ent["label"] in ENTITY_TYPES | |
| assert ent["start_token"] <= ent["end_token"] | |
| assert 0.0 <= ent["confidence"] <= 1.0 | |
| def test_entity_span_text_matches_tokens(self): | |
| from finner.infer.predict import _clean_span, predict | |
| result = predict("Goldman Sachs earned $500M in Q3 2024.") | |
| tokens = result["tokens"] | |
| for ent in result["entities"]: | |
| raw = " ".join(tokens[ent["start_token"]: ent["end_token"] + 1]) | |
| # entity text is the cleaned (trailing-punct + unbalanced-paren stripped) span | |
| assert _clean_span(raw) == ent["text"] | |
| def test_empty_text_edge_case(self): | |
| from finner.infer.predict import predict | |
| result = predict(".") | |
| assert isinstance(result["entities"], list) | |