File size: 10,402 Bytes
5e952bb | 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 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 | import json
import logging
from src.config import settings
from src.models.claim import Claim
from src.llm import get_llm
from pydantic import BaseModel, Field
logger = logging.getLogger(__name__)
# Try to detect if scispacy and its dependencies are available (compiled / installed)
SCISPACY_AVAILABLE = False
try:
import spacy # noqa: F401
import scispacy # noqa: F401
from scispacy.linking import EntityLinker # noqa: F401
SCISPACY_AVAILABLE = True
except ImportError:
pass
# Default synonym and canonical concept dictionary mapping (case-insensitive keys)
# Fast lookup for common synonyms to avoid calling external models
DEFAULT_SYNONYM_MAP = {
"aspirin": ("Aspirin", "MeSH:D001241"),
"asa": ("Aspirin", "MeSH:D001241"),
"acetylsalicylic acid": ("Aspirin", "MeSH:D001241"),
"metformin": ("Metformin", "MeSH:D008687"),
"glucophage": ("Metformin", "MeSH:D008687"),
"cancer": ("Neoplasms", "MeSH:D009369"),
"tumor": ("Neoplasms", "MeSH:D009369"),
"tumors": ("Neoplasms", "MeSH:D009369"),
"neoplasm": ("Neoplasms", "MeSH:D009369"),
"neoplasms": ("Neoplasms", "MeSH:D009369"),
"breast cancer": ("Breast Neoplasms", "MeSH:D001943"),
"breast tumors": ("Breast Neoplasms", "MeSH:D001943"),
"ampk": ("AMP-activated Protein Kinases", "MeSH:D055372"),
"mtor": ("TOR Serine-Threonine Kinases", "MeSH:D058570"),
}
class NormalizedEntity(BaseModel):
"""Pydantic model representing a single normalized entity resolved by the LLM."""
original_text: str = Field(description="The original entity text mention to match back to")
canonical_name: str = Field(description="The normalized canonical name of the entity concept")
canonical_id: str | None = Field(default=None, description="Canonical database ID, e.g., 'MeSH:D001241' or 'UMLS:C0004057'")
entity_type: str = Field(description="The biomedical entity type, e.g., 'DRUG', 'DISEASE', etc.")
class EntityNormalizationResponse(BaseModel):
"""Schema for the LLM structured response in entity normalization."""
normalized_entities: list[NormalizedEntity]
class EntityNormalizer:
def __init__(self, synonym_map_path: str | None = None):
self.nlp = None
self._scispacy_initialized = False
self.synonym_map = self._load_synonym_map(synonym_map_path)
def _load_synonym_map(self, synonym_map_path: str | None = None) -> dict[str, tuple[str, str | None]]:
"""Load the local synonym map from a JSON file, falling back to default map on error."""
if synonym_map_path is None:
synonym_map_path = settings.synonym_map_path
if not synonym_map_path:
logger.info("No synonym map path configured. Using default synonym map.")
return DEFAULT_SYNONYM_MAP
import os
resolved_path = os.path.abspath(synonym_map_path)
if not os.path.exists(resolved_path):
logger.warning(f"Synonym map file not found at '{resolved_path}'. Falling back to default synonym map.")
return DEFAULT_SYNONYM_MAP
try:
with open(resolved_path, "r", encoding="utf-8") as f:
data = json.load(f)
if not isinstance(data, dict):
logger.warning(f"Synonym map file at '{resolved_path}' must contain a JSON object. Falling back to default synonym map.")
return DEFAULT_SYNONYM_MAP
loaded_map = {}
for k, v in data.items():
if isinstance(v, list) and len(v) >= 2:
loaded_map[k.lower()] = (v[0], v[1])
elif isinstance(v, list) and len(v) == 1:
loaded_map[k.lower()] = (v[0], None)
else:
logger.warning(f"Invalid value for synonym key '{k}': {v}. Expected a list of [canonical_name, canonical_id]. Skipping.")
logger.info(f"Loaded {len(loaded_map)} synonyms from '{resolved_path}'.")
return loaded_map
except Exception as e:
logger.warning(f"Failed to load synonym map from '{resolved_path}': {e}. Falling back to default synonym map.")
return DEFAULT_SYNONYM_MAP
def _init_scispacy(self) -> None:
"""Lazily initialize scispaCy pipeline to avoid import delays if unused."""
if not SCISPACY_AVAILABLE or self._scispacy_initialized:
return
try:
import spacy
logger.info("Initializing scispaCy NLP pipeline and UMLS EntityLinker...")
self.nlp = spacy.load("en_core_sci_sm")
self.nlp.add_pipe(
"scispacy_linker",
config={"resolve_abbreviations": True, "linker_name": "umls"}
)
self._scispacy_initialized = True
logger.info("scispaCy pipeline successfully initialized.")
except Exception as e:
logger.warning(f"Failed to initialize scispaCy pipeline: {e}. Falling back to LLM / local map.")
self._scispacy_initialized = False
async def normalize_entities(self, claims: list[Claim]) -> list[Claim]:
"""Normalize the entities in a list of Claims.
Updates the claims' entities in-place by assigning canonical names and IDs.
Resolves synonyms using a local dictionary map, scispaCy (if available),
and falls back to LLM-based structured resolution.
"""
# 1. Gather all unique entity mentions across claims to process in batch
unique_entities = {}
for claim in claims:
for entity in claim.entities:
text_clean = entity.text.strip()
if not text_clean:
continue
key = text_clean.lower()
if key not in unique_entities:
unique_entities[key] = {
"text": text_clean,
"type": entity.entity_type
}
if not unique_entities:
return claims
resolved_map = {} # key (lowercase original text) -> (canonical_name, canonical_id)
unresolved_entities = []
# 2. Check local fast-lookup synonym map
for key, info in unique_entities.items():
if key in self.synonym_map:
resolved_map[key] = self.synonym_map[key]
else:
unresolved_entities.append(info)
# 3. If scispaCy is available, attempt to resolve using the linker
if unresolved_entities and SCISPACY_AVAILABLE:
self._init_scispacy()
if self._scispacy_initialized:
still_unresolved = []
for info in unresolved_entities:
text = info["text"]
try:
doc = self.nlp(text)
if doc.ents and doc.ents[0]._.kb_ents:
umls_ent_id, _ = doc.ents[0]._.kb_ents[0]
kb = self.nlp.get_pipe("scispacy_linker").kb
canonical_name = kb.cui_to_entity[umls_ent_id].canonical_name
canonical_id = f"UMLS:{umls_ent_id}"
resolved_map[text.lower()] = (canonical_name, canonical_id)
else:
still_unresolved.append(info)
except Exception as e:
logger.warning(f"scispaCy linking failed for '{text}': {e}")
still_unresolved.append(info)
unresolved_entities = still_unresolved
# 4. Fallback to LLM structured resolution for remaining unresolved concepts
if unresolved_entities:
try:
llm = get_llm()
# Prepare payload for LLM normalizer
entities_to_resolve = [
{"text": info["text"], "type": info["type"].value}
for info in unresolved_entities
]
logger.info(f"Resolving {len(entities_to_resolve)} entities via LLM ({llm.model_name})...")
prompt = (
"You are a medical informatics expert. Normalize the following biomedical entity mentions to their canonical names "
"and standard database identifiers (preferably MeSH descriptors, e.g., 'MeSH:D001241', or UMLS CUIs, e.g., 'UMLS:C0004057').\n\n"
"Rules:\n"
"1. Identify the most specific canonical term/concept name for the given text.\n"
"2. Provide the canonical identifier in the format 'MeSH:ID' or 'UMLS:ID'. If no standard ID exists, return null for canonical_id.\n"
"3. Keep the original text exactly as provided to match back to the input.\n"
"4. Resolve synonyms to the same canonical concept. For example, 'ASA', 'acetylsalicylic acid', and 'aspirin' must all map to "
"canonical name 'Aspirin' and canonical ID 'MeSH:D001241'.\n\n"
f"Entities to normalize:\n{json.dumps(entities_to_resolve, indent=2)}\n\n"
"Please respond with the normalized entities in JSON format matching the schema."
)
response = await llm.generate_structured(
prompt=prompt,
response_schema=EntityNormalizationResponse,
temperature=0.1
)
for ent in response.normalized_entities:
key = ent.original_text.strip().lower()
resolved_map[key] = (ent.canonical_name, ent.canonical_id)
except Exception as e:
logger.error(f"LLM entity normalization failed: {e}. Keeping original names.")
# 5. Apply the resolved mapping to update claims' entities in-place
for claim in claims:
for entity in claim.entities:
key = entity.text.strip().lower()
if key in resolved_map:
canonical_name, canonical_id = resolved_map[key]
entity.text = canonical_name
entity.canonical_id = canonical_id
return claims
|