Hugging Face's logo

Spaces:
Lp012
/
RSCE
Sleeping

App Files Community
RSCE / tests /test_normalizer.py
LP1205's picture
LP1205
Deploy RSCE backend to Hugging Face Space
5e952bb
Raw History Blame Contribute Delete
6.56 kB
import pytest
import uuid
import json
from unittest.mock import AsyncMock, MagicMock, patch
from src.models.claim import Claim, ClaimType, Polarity, StudyDesign, Entity, EntityType
from src.entity.normalizer import EntityNormalizer, EntityNormalizationResponse, NormalizedEntity
@pytest.fixture
def sample_claims():
entity_asa = Entity(text="ASA", entity_type=EntityType.DRUG)
entity_acid = Entity(text="acetylsalicylic acid", entity_type=EntityType.DRUG)
entity_aspirin = Entity(text="aspirin", entity_type=EntityType.DRUG)
entity_metformin = Entity(text="glucophage", entity_type=EntityType.DRUG)
entity_insulin = Entity(text="insulin", entity_type=EntityType.DRUG) # not in local map, will trigger LLM fallback
claim_1 = Claim(
id=uuid.uuid4(),
text="ASA reduces headache.",
paper_id="11111",
year=2020,
confidence_score=1.0,
claim_type=ClaimType.CAUSAL,
polarity=Polarity.NEGATIVE,
entities=[entity_asa, entity_acid],
population="humans",
context="general",
quote_anchor="reduces headache",
study_design=StudyDesign.RCT
)
claim_2 = Claim(
id=uuid.uuid4(),
text="Aspirin and glucophage treatment.",
paper_id="22222",
year=2023,
confidence_score=1.0,
claim_type=ClaimType.CAUSAL,
polarity=Polarity.POSITIVE,
entities=[entity_aspirin, entity_metformin, entity_insulin],
population="humans",
context="general",
quote_anchor="treatment",
study_design=StudyDesign.RCT
)
return [claim_1, claim_2]
@pytest.mark.asyncio
async def test_normalize_entities_local_lookup(sample_claims):
# Ensure that local lookup works and doesn't hit the LLM since we don't mock it here
# (except for 'insulin', which we should mock or test separately. So let's temporarily mock the LLM
# call to return a mock response for 'insulin')
mock_llm = MagicMock()
mock_llm.model_name = "mock-llm"
mock_llm.generate_structured = AsyncMock(
return_value=EntityNormalizationResponse(
normalized_entities=[
NormalizedEntity(
original_text="insulin",
canonical_name="Insulin",
canonical_id="MeSH:D007328",
entity_type="DRUG"
)
]
)
)
with patch("src.entity.normalizer.get_llm", return_value=mock_llm):
normalizer = EntityNormalizer()
normalized_claims = await normalizer.normalize_entities(sample_claims)
# Verify claim 1 entities (ASA and acetylsalicylic acid) both resolved to Aspirin / MeSH:D001241
assert normalized_claims[0].entities[0].text == "Aspirin"
assert normalized_claims[0].entities[0].canonical_id == "MeSH:D001241"
assert normalized_claims[0].entities[1].text == "Aspirin"
assert normalized_claims[0].entities[1].canonical_id == "MeSH:D001241"
# Verify claim 2 entities
# aspirin -> Aspirin / MeSH:D001241
assert normalized_claims[1].entities[0].text == "Aspirin"
assert normalized_claims[1].entities[0].canonical_id == "MeSH:D001241"
# glucophage -> Metformin / MeSH:D008687
assert normalized_claims[1].entities[1].text == "Metformin"
assert normalized_claims[1].entities[1].canonical_id == "MeSH:D008687"
# insulin -> Insulin / MeSH:D007328 (from LLM)
assert normalized_claims[1].entities[2].text == "Insulin"
assert normalized_claims[1].entities[2].canonical_id == "MeSH:D007328"
# Verify LLM was called exactly once for the single unresolved entity 'insulin'
mock_llm.generate_structured.assert_called_once()
call_kwargs = mock_llm.generate_structured.call_args[1]
assert "insulin" in call_kwargs["prompt"]
# And fast-lookup entities were NOT sent to the LLM as input
assert "glucophage" not in call_kwargs["prompt"]
@pytest.mark.asyncio
async def test_normalize_entities_llm_fallback_failure(sample_claims):
# Test that if the LLM fallback fails, original names are kept for unresolved entities
mock_llm = MagicMock()
mock_llm.model_name = "mock-llm"
mock_llm.generate_structured = AsyncMock(side_effect=Exception("API failure"))
with patch("src.entity.normalizer.get_llm", return_value=mock_llm):
normalizer = EntityNormalizer()
normalized_claims = await normalizer.normalize_entities(sample_claims)
# Local lookup ones are still resolved
assert normalized_claims[0].entities[0].text == "Aspirin"
assert normalized_claims[1].entities[1].text == "Metformin"
# Insulin (which failed LLM resolution) kept its original text and None canonical_id
assert normalized_claims[1].entities[2].text == "insulin"
assert normalized_claims[1].entities[2].canonical_id is None
@pytest.mark.asyncio
async def test_normalizer_loads_custom_synonym_map(tmp_path):
# Create a custom synonym map file
custom_map = {
"advil": ["Ibuprofen", "MeSH:D007052"],
"tylenol": ["Acetaminophen", "MeSH:D000082"]
}
custom_file = tmp_path / "custom_synonyms.json"
custom_file.write_text(json.dumps(custom_map))
normalizer = EntityNormalizer(synonym_map_path=str(custom_file))
# Assert custom map loaded correctly
assert normalizer.synonym_map["advil"] == ("Ibuprofen", "MeSH:D007052")
assert normalizer.synonym_map["tylenol"] == ("Acetaminophen", "MeSH:D000082")
# Assert fallback default items are NOT in custom map (since it loaded custom only)
assert "aspirin" not in normalizer.synonym_map
def test_normalizer_fallback_on_missing_file():
# Instantiate with non-existent path
normalizer = EntityNormalizer(synonym_map_path="non_existent_file.json")
# Should fall back to default map
assert "aspirin" in normalizer.synonym_map
assert normalizer.synonym_map["aspirin"] == ("Aspirin", "MeSH:D001241")
def test_normalizer_fallback_on_invalid_json(tmp_path):
# Create an invalid JSON file
invalid_file = tmp_path / "invalid.json"
invalid_file.write_text("invalid json content {")
normalizer = EntityNormalizer(synonym_map_path=str(invalid_file))
# Should fall back to default map
assert "aspirin" in normalizer.synonym_map
assert normalizer.synonym_map["aspirin"] == ("Aspirin", "MeSH:D001241")