|
|
|
|
|
import spacy |
|
|
from scispacy.linking import EntityLinker |
|
|
|
|
|
@spacy.util.cache_dir("~/.cache/scispacy") |
|
|
def load_model(): |
|
|
nlp = spacy.load("en_core_sci_scibert") |
|
|
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) -> list[dict]: |
|
|
""" |
|
|
Returns unique UMLS concepts with confidence scores and semantic types. |
|
|
""" |
|
|
doc = nlp(text) |
|
|
best = {} |
|
|
for ent in doc.ents: |
|
|
for cui, score in ent._.umls_ents: |
|
|
meta = nlp.get_pipe("scispacy_linker").kb.cui_to_entity[cui] |
|
|
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()) |
|
|
|
