A neural-symbolic AI agent system for biomedical concept mapping.

1. from collections import Counter
2. from collections import defaultdict
3. from langchain_ollama import ChatOllama
4. from mcp.server.fastmcp import FastMCP
5. from pydantic import BaseModel, ValidationError
6. from scipy.sparse import save_npz, load_npz
7. from scispacy.linking import CandidateGenerator
8. from sklearn.feature_extraction.text import TfidfVectorizer
9. from sklearn.metrics.pairwise import cosine_similarity
10. from typing import Any, List, Dict, Optional, Union
11. from typing import Type, TypeVar
12. from utils import umls
14. import hydra
15. import json
16. import logging
17. import numpy as np
18. import pickle
20. with hydra.initialize(config_path="../config", version_base=None):
21. similary_context_config = hydra.compose(config_name="ann")
22. similary_context_ann = json.load(open(similary_context_config.data_path))
23. similary_context_encoding = pickle.load(open(similary_context_config.encoding_path, "rb"))
24. similary_context_vectorizer = pickle.load(open(similary_context_config.vec_path, "rb"))
25. similary_context_threshold = similary_context_config.threshold
27. with hydra.initialize(config_path="../config", version_base=None):
28. ollama_config = hydra.compose(config_name="ollama")
29. llm = ChatOllama(
30. model=ollama_config.model,
31. base_url="http://localhost:11434",
32. # base_url=ollama_config.base_url,
33. )
35. with hydra.initialize(config_path="../config", version_base=None):
36. umls_config = hydra.compose(config_name="umls")
37. umls_cuis = json.load(open(umls_config.cui_path))
38. umls_encoding = pickle.load(open(umls_config.encoding_path, "rb"))
39. umls_vectorizer = pickle.load(open(umls_config.vec_path, "rb"))
40. umls_threshold = umls_config.threshold
41. umls_top_k = umls_config.top_k
43. cui_to_names = umls.load_umls_names_and_aliases()
46. def extract_json_block(text: str) -> str:
47. try:
48. start = text.index('{')
49. end = text.rindex('}') + 1
50. return text[start:end].strip()
51. except ValueError:
52. raise ValueError(f"No valid JSON object found in the input string: {text}.")
54. T = TypeVar("T", bound=BaseModel)
55. def parse_llm_response_as_model(raw_response: str, model: Type[T]) -> T:
56. """
57. Parse and validate a JSON string from LLM output using the given Pydantic model.
58. """
59. try:
60. content = raw_response.strip()
61. content = extract_json_block(content)
62. content = content.replace("\\n", "\n").replace("\\'", "'")
63. parsed_json = json.loads(content)
64. return model(**parsed_json)
65. except (json.JSONDecodeError, ValidationError) as e:
66. raise ValueError(f"Parsing failed: {e}")
68. class EquivalenceJudgment(BaseModel):
69. result: bool
70. explanation: str
72. def judge_refinement_equivalence(
73. mention: str,
74. context_snippet: str,
75. refined_phrase: str
76. ) -> Dict[str, Union[bool, str]]:
77. """
78. Return a dict containing whether the refined phrase is semantically equivalent to the original mention,
79. and an explanation.
80. """
81. prompt = (
82. f"You are a biomedical expert evaluating whether a refined phrase reasonably captures the same meaning "
83. f"as the original concept mention, given the surrounding clinical context.\n\n"
84. f"Original mention: \"{mention}\"\n"
85. f"Context: \"{context_snippet}\"\n"
86. f"Refined phrase: \"{refined_phrase}\"\n\n"
87. f"Only respond `true` if the refined phrase can reasonably be interpreted to mean the same thing "
88. f"in context, even if there are minor wording differences.\n"
89. f"Only respond `false` if the meaning is clearly different.\n\n"
90. f"Respond in JSON format:\n"
91. '{"result": true/false, "explanation": "brief justification"}'
92. )
94. try:
95. response = llm.invoke(prompt)
96. return parse_llm_response_as_model(response.content, EquivalenceJudgment).__dict__
97. except Exception as e:
98. return {
99. "result": False,
100. "explanation": f"Fallback due to error: {str(e)}"
101. }
103. def search_similar_contexts(concept: str) -> dict:
104. """Search contexts that are similar to the input."""
105. vector = similary_context_vectorizer.transform([concept])
106. similarity_matrix = cosine_similarity(vector, similary_context_encoding)
107. idx_to_examples = {}
108. similarities = similarity_matrix[0]
109. matching_indices = np.where(similarities >= similary_context_threshold)[0]
110. matching_scores = similarities[matching_indices]
111. examples = []
112. cuis = []
113. for candidate_idx in matching_indices:
114. candidate = similary_context_ann[candidate_idx]
115. cui = candidate['cuis'][0]
116. cuis.append(cui)
117. examples.append(candidate)
119. similar_context = {
120. 'input': concept,
121. 'similar_contexts': examples,
122. }
123. return similar_context
125. def search_nearest_neighbors(concept: str, context_snippet: str) -> List[str]:
126. """
127. Retrieve CUIs that appear in similar contexts and whose representative names are
128. semantically equivalent to the input concept within the provided context.
130. Args:
131. concept: The original concept mention.
132. context_snippet: Context around the concept (e.g., sentence or paragraph).
134. Returns:
135. A list of CUIs that are semantically equivalent based on neighboring context.
136. """
137. similar_context = search_similar_contexts(concept)
138. candidate_cuis = []
140. for context in similar_context.get('similar_contexts', []):
141. candidate_cuis.extend(context.get('cuis', []))
143. cuis = []
144. cui_counter = Counter(candidate_cuis)
145. return [cui for cui, _ in cui_counter.most_common()]
148. class AugmentationJudgment(BaseModel):
149. result: bool
150. explanation: str
153. def judge_need_for_augmentation(
154. mention: str,
155. context_snippet: str
156. ) -> AugmentationJudgment:
157. prompt = (
158. f"You are a biomedical language expert. Given the following concept mention, "
159. f"determine whether the mention is expressed as regular words.\n\n"
160. f"Concept mention: \"{mention}\"\n"
161. f"If the term is not a regular word, respond `true`.\n"
162. f"Otherwise, respond `false`.\n\n"
163. f"Respond in JSON format:\n"
164. f"{{\"result\": true/false, \"explanation\": \"brief justification\"}}"
165. )
167. try:
168. response = llm.invoke(prompt)
169. return parse_llm_response_as_model(response.content, AugmentationJudgment).__dict__
170. except Exception as e:
171. return {
172. "result": False,
173. "explanation": f"Fallback due to error: {str(e)}"
174. }
176. class AugmentedConcept(BaseModel):
177. augmented_mention: str
178. explanation: str
180. def augment_concept_mention(
181. mention: str,
182. context_snippet: str
183. ) -> Dict[str, Union[str, bool]]:
184. """
185. Return a dict containing an augmented version of the mention (if applicable) and an explanation.
186. If the original mention is already clear and standard, the augmented form may be identical.
187. """
188. prompt = (
189. f"You are a biomedical expert tasked with improving the clarity of a clinical concept mention.\n\n"
190. f"Concept mention: \"{mention}\"\n"
191. f"Context: \"{context_snippet}\"\n\n"
192. f"If the mention is abbreviated or shortened, replace it with a fully spelled term mentioned in the context."
193. f"If the mention is already standard and clear, keep it unchanged.\n\n"
194. f"Respond in JSON format:\n"
195. '{"augmented_mention": "string", "explanation": "brief justification"}'
196. )
197. try:
198. response = llm.invoke(prompt)
199. return parse_llm_response_as_model(response.content, AugmentedConcept).__dict__
200. except Exception as e:
201. return {
202. "augmented_mention": mention,
203. "explanation": f"Fallback due to error: {str(e)}. Response: {response.content}"
204. }
206. def link_mention_to_concept_ids(concept: str) -> List[str]:
207. """
208. Given a concept mention (text string), return a ranked list of UMLS concept IDs
209. that are semantically similar to the input mention.
210. """
211. concept_ids = []
212. vector = umls_vectorizer.transform([concept])
213. similarity_matrix = cosine_similarity(vector, umls_encoding)
214. similarities = similarity_matrix[0]
215. matching_indices = np.where(similarities >= umls_threshold)[0]
216. ranked_indices = matching_indices[np.argsort(
217. similarities[matching_indices])[::-1]]
218. for candidate_idx in ranked_indices:
219. cui = umls_cuis[candidate_idx]
220. if not cui in concept_ids:
221. concept_ids.append(cui)
222. if len(concept_ids) >= umls_top_k:
223. break
224. return concept_ids
226. def link_mentions_rrf(candidates: List[str]) -> List[str]:
227. """
228. Given multiple textual variants of the same concept, return a merged ranked list of UMLS CUIs
229. using Reciprocal Rank Fusion (RRF) with batched TF-IDF vectorization.
231. Returns:
232. A ranked list of UMLS CUIs.
233. """
234. k, rrf_k = umls_top_k, 60
235. fusion_scores: Dict[str, float] = defaultdict(float)
237. # Batch vectorize all candidates
238. vectors = umls_vectorizer.transform(candidates) # shape: (n_candidates, vocab_dim)
239. similarity_matrix = cosine_similarity(vectors, umls_encoding) # shape: (n_candidates, n_umls)
241. for i, name in enumerate(candidates):
242. similarities = similarity_matrix[i]
243. matching_indices = np.where(similarities >= umls_threshold)[0]
244. ranked_indices = matching_indices[np.argsort(similarities[matching_indices])[::-1]]
246. for rank, idx in enumerate(ranked_indices):
247. cui = umls_cuis[idx]
248. fusion_scores[cui] += 1.0 / (rrf_k + rank + 1)
249. if rank + 1 >= umls_top_k:
250. break
252. # Sort by descending RRF score
253. sorted_cuis = sorted(fusion_scores.items(), key=lambda x: x[1], reverse=True)
254. return [cui for cui, _ in sorted_cuis[:k]]
259. class ConceptExtraction(BaseModel):
260. concepts: List[str]
261. explanation: str
263. def extract_biomedical_concepts(passage: str) -> Dict[str, Union[List[str], str]]:
264. """
265. Extract biomedical concepts from a clinical sentence or paragraph.
266. """
267. prompt = (
268. "You are a biomedical NLP assistant. Given the following clinical sentence or paragraph, "
269. "identify and list key PICO-related concepts. These may include population groups, interventions, "
270. "comparators, and outcomes. Only include medically relevant phrases. "
271. "If no relevant concepts are found, return an empty list.\n\n"
272. f"Text:\n\"{passage}\"\n\n"
273. "Respond in the following JSON format:\n"
274. "{\n"
275. " \"concepts\": [\"...\", \"...\", ...]\n"
276. "}"
277. )
279. try:
280. response = llm.invoke(prompt)
281. return parse_llm_response_as_model(response.content, ConceptExtraction).__dict__
283. except Exception as e:
284. return {"concepts": [], "explanation": f"Error: {str(e)}"}

concept.py at commit 0e996f5, no license · at the source