Update mcp/nlp.py

mcp/nlp.py

@@ -1,38 +1,30 @@
 # mcp/nlp.py
 import spacy
-import scispacy
 from scispacy.linking import EntityLinker
 
-# Load a powerful biomedical model + UMLS linker
 @spacy.util.cache_dir("~/.cache/scispacy")
 def load_model():
     nlp = spacy.load("en_core_sci_scibert")
-    # Resolve abbreviations then link to UMLS
     linker = EntityLinker(name="umls", resolve_abbreviations=True, threshold=0.75)
     nlp.add_pipe(linker)
     return nlp
 
 nlp = load_model()
 
-def extract_umls_concepts(text: str):
+def extract_umls_concepts(text: str) -> list[dict]:
     """
-    Returns a list of {cui, concept_name, score, semantic_types}.
+    Returns unique UMLS concepts with confidence scores and semantic types.
     """
     doc = nlp(text)
-    concepts = []
+    best = {}
     for ent in doc.ents:
         for cui, score in ent._.umls_ents:
             meta = nlp.get_pipe("scispacy_linker").kb.cui_to_entity[cui]
-            concepts.append({
-                "cui": cui,
-                "name": meta.canonical_name,
-                "score": float(score),
-                "types": meta.types  # list of semantic type strings
-            })
-    # Deduplicate by CUI, keep highest score
-    seen = {}
-    for c in concepts:
-        prev = seen.get(c["cui"])
-        if not prev or c["score"] > prev["score"]:
-            seen[c["cui"]] = c
-    return list(seen.values())
+            if cui not in best or score > best[cui]["score"]:
+                best[cui] = {
+                    "cui": cui,
+                    "name": meta.canonical_name,
+                    "score": float(score),
+                    "types": meta.types
+                }
+    return list(best.values())