Sentence Entity Resolver for SNOMED CT (Body Structures) (bge_base_en_v1_5_onnx embeddings)

Description

This model maps extracted clinical NER entities to SNOMED CT concepts using bge_base_en_v1_5_onnx embeddings.

It is trained on SNOMED CT US Edition 20260901 release.

Live Demo Open in Colab Copy S3 URI

How to use

documentAssembler = DocumentAssembler()\
    .setInputCol("text")\
    .setOutputCol("document")

sentenceDetectorDL = SentenceDetectorDLModel.pretrained("sentence_detector_dl_healthcare", "en", "clinical/models")\
    .setInputCols(["document"])\
    .setOutputCol("sentence")

tokenizer = Tokenizer()\
    .setInputCols(["sentence"])\
    .setOutputCol("token")

word_embeddings = WordEmbeddingsModel.pretrained("embeddings_clinical","en","clinical/models")\
    .setInputCols(["sentence","token"])\
    .setOutputCol("embeddings")

ner_jsl = MedicalNerModel.pretrained("ner_jsl","en","clinical/models")\
    .setInputCols(["sentence","token","embeddings"])\
    .setOutputCol("ner_jsl_tags")

ner_jsl_converter = NerConverterInternal()\
    .setInputCols(["sentence","token","ner_jsl_tags"])\
    .setOutputCol("ner_chunk_jsl")\
    .setWhiteList(["Disease_Syndrome_Disorder", "External_body_part_or_region"])\
    .setReplaceLabels({"Disease_Syndrome_Disorder": "BodyPart", "External_body_part_or_region": "BodyPart"})

ner_anatomy = MedicalNerModel.pretrained("ner_anatomy_coarse","en","clinical/models")\
    .setInputCols(["sentence","token","embeddings"])\
    .setOutputCol("ner_anatomy_tags")

ner_anatomy_converter = NerConverterInternal()\
    .setInputCols(["sentence","token","ner_anatomy_tags"])\
    .setOutputCol("ner_chunk_anatomy")\
    .setReplaceLabels({"Anatomy": "BodyPart"})

ner_oncology_anatomy = MedicalNerModel.pretrained("ner_oncology_anatomy_general","en","clinical/models")\
    .setInputCols(["sentence","token","embeddings"])\
    .setOutputCol("ner_oncology_anatomy_tags")

ner_oncology_anatomy_converter = NerConverterInternal()\
    .setInputCols(["sentence","token","ner_oncology_anatomy_tags"])\
    .setOutputCol("ner_chunk_oncology_anatomy")\
    .setReplaceLabels({"Anatomical_Site": "BodyPart"})

chunk_merger = ChunkMergeApproach()\
    .setInputCols(["ner_chunk_jsl", "ner_chunk_anatomy", "ner_chunk_oncology_anatomy"])\
    .setOutputCol("ner_chunk")

chunk2doc = Chunk2Doc()\
    .setInputCols(["ner_chunk"])\
    .setOutputCol("ner_chunk_doc")

embedder = BGEEmbeddings.pretrained("bge_base_en_v1_5_onnx", "en")\
    .setInputCols(["ner_chunk_doc"])\
    .setOutputCol("bge_embeddings")\
    .setCaseSensitive(False)

resolver = SentenceEntityResolverModel.pretrained("bgeresolve_snomed_bodyStructure_20260901","en","clinical/models")\
    .setInputCols(["bge_embeddings"])\
    .setOutputCol("snomed_code")\
    .setDistanceFunction("EUCLIDEAN")\
    .setThreshold(1000)

pipeline = Pipeline(stages=[\
    documentAssembler, sentenceDetectorDL, tokenizer, word_embeddings, ner_jsl, ner_jsl_converter, ner_anatomy, ner_anatomy_converter, ner_oncology_anatomy, ner_oncology_anatomy_converter, chunk_merger, chunk2doc, embedder, resolver\
])

data = spark.createDataFrame([["The patient is a 30-year-old female with coronary artery disease and swelling affecting the kidney and lower limb."]]).toDF("text")
result = pipeline.fit(data).transform(data)
documentAssembler = nlp.DocumentAssembler()\
    .setInputCol("text")\
    .setOutputCol("document")

sentenceDetectorDL = nlp.SentenceDetectorDLModel.pretrained("sentence_detector_dl_healthcare", "en", "clinical/models")\
    .setInputCols(["document"])\
    .setOutputCol("sentence")

tokenizer = nlp.Tokenizer()\
    .setInputCols(["sentence"])\
    .setOutputCol("token")

word_embeddings = nlp.WordEmbeddingsModel.pretrained("embeddings_clinical","en","clinical/models")\
    .setInputCols(["sentence","token"])\
    .setOutputCol("embeddings")

ner_jsl = medical.NerModel.pretrained("ner_jsl","en","clinical/models")\
    .setInputCols(["sentence","token","embeddings"])\
    .setOutputCol("ner_jsl_tags")

ner_jsl_converter = medical.NerConverterInternal()\
    .setInputCols(["sentence","token","ner_jsl_tags"])\
    .setOutputCol("ner_chunk_jsl")\
    .setWhiteList(["Disease_Syndrome_Disorder", "External_body_part_or_region"])\
    .setReplaceLabels({"Disease_Syndrome_Disorder": "BodyPart", "External_body_part_or_region": "BodyPart"})

ner_anatomy = medical.NerModel.pretrained("ner_anatomy_coarse","en","clinical/models")\
    .setInputCols(["sentence","token","embeddings"])\
    .setOutputCol("ner_anatomy_tags")

ner_anatomy_converter = medical.NerConverterInternal()\
    .setInputCols(["sentence","token","ner_anatomy_tags"])\
    .setOutputCol("ner_chunk_anatomy")\
    .setReplaceLabels({"Anatomy": "BodyPart"})

ner_oncology_anatomy = medical.NerModel.pretrained("ner_oncology_anatomy_general","en","clinical/models")\
    .setInputCols(["sentence","token","embeddings"])\
    .setOutputCol("ner_oncology_anatomy_tags")

ner_oncology_anatomy_converter = medical.NerConverterInternal()\
    .setInputCols(["sentence","token","ner_oncology_anatomy_tags"])\
    .setOutputCol("ner_chunk_oncology_anatomy")\
    .setReplaceLabels({"Anatomical_Site": "BodyPart"})

chunk_merger = medical.ChunkMergeApproach()\
    .setInputCols(["ner_chunk_jsl", "ner_chunk_anatomy", "ner_chunk_oncology_anatomy"])\
    .setOutputCol("ner_chunk")

chunk2doc = nlp.Chunk2Doc()\
    .setInputCols(["ner_chunk"])\
    .setOutputCol("ner_chunk_doc")

embedder = nlp.BGEEmbeddings.pretrained("bge_base_en_v1_5_onnx", "en")\
    .setInputCols(["ner_chunk_doc"])\
    .setOutputCol("bge_embeddings")\
    .setCaseSensitive(False)

resolver = medical.SentenceEntityResolverModel.pretrained("bgeresolve_snomed_bodyStructure_20260901","en","clinical/models")\
    .setInputCols(["bge_embeddings"])\
    .setOutputCol("snomed_code")\
    .setDistanceFunction("EUCLIDEAN")\
    .setThreshold(1000)

pipeline = nlp.Pipeline(stages=[\
    documentAssembler, sentenceDetectorDL, tokenizer, word_embeddings, ner_jsl, ner_jsl_converter, ner_anatomy, ner_anatomy_converter, ner_oncology_anatomy, ner_oncology_anatomy_converter, chunk_merger, chunk2doc, embedder, resolver\
])

data = spark.createDataFrame([["The patient is a 30-year-old female with coronary artery disease and swelling affecting the kidney and lower limb."]]).toDF("text")
result = pipeline.fit(data).transform(data)

val documentAssembler = new DocumentAssembler()
    .setInputCol("text")
    .setOutputCol("document")

val sentenceDetectorDL = SentenceDetectorDLModel
    .pretrained("sentence_detector_dl_healthcare", "en", "clinical/models")
    .setInputCols(Array("document"))
    .setOutputCol("sentence")

val tokenizer = new Tokenizer()
    .setInputCols("sentence")
    .setOutputCol("token")

val word_embeddings = WordEmbeddingsModel
    .pretrained("embeddings_clinical", "en", "clinical/models")
    .setInputCols(Array("sentence", "token"))
    .setOutputCol("embeddings")

val ner_jsl = MedicalNerModel
    .pretrained("ner_jsl", "en", "clinical/models")
    .setInputCols(Array("sentence", "token", "embeddings"))
    .setOutputCol("ner_jsl_tags")

val ner_jsl_converter = new NerConverterInternal()
    .setInputCols(Array("sentence", "token", "ner_jsl_tags"))
    .setOutputCol("ner_chunk_jsl")
    .setWhiteList(Array("Disease_Syndrome_Disorder", "External_body_part_or_region"))
    .setReplaceLabels(Map("Disease_Syndrome_Disorder" -> "BodyPart", "External_body_part_or_region" -> "BodyPart"))

val ner_anatomy = MedicalNerModel
    .pretrained("ner_anatomy_coarse", "en", "clinical/models")
    .setInputCols(Array("sentence", "token", "embeddings"))
    .setOutputCol("ner_anatomy_tags")

val ner_anatomy_converter = new NerConverterInternal()
    .setInputCols(Array("sentence", "token", "ner_anatomy_tags"))
    .setOutputCol("ner_chunk_anatomy")
    .setReplaceLabels(Map("Anatomy" -> "BodyPart"))

val ner_oncology_anatomy = MedicalNerModel
    .pretrained("ner_oncology_anatomy_general", "en", "clinical/models")
    .setInputCols(Array("sentence", "token", "embeddings"))
    .setOutputCol("ner_oncology_anatomy_tags")

val ner_oncology_anatomy_converter = new NerConverterInternal()
    .setInputCols(Array("sentence", "token", "ner_oncology_anatomy_tags"))
    .setOutputCol("ner_chunk_oncology_anatomy")
    .setReplaceLabels(Map("Anatomical_Site" -> "BodyPart"))

val chunk_merger = new ChunkMergeApproach()
    .setInputCols(Array("ner_chunk_jsl", "ner_chunk_anatomy", "ner_chunk_oncology_anatomy"))
    .setOutputCol("ner_chunk")

val chunk2doc = new Chunk2Doc()
    .setInputCols(Array("ner_chunk"))
    .setOutputCol("ner_chunk_doc")

val embedder = BGEEmbeddings
    .pretrained("bge_base_en_v1_5_onnx", "en")
    .setInputCols(Array("ner_chunk_doc"))
    .setOutputCol("bge_embeddings")
    .setCaseSensitive(false)

val resolver = SentenceEntityResolverModel
    .pretrained("bgeresolve_snomed_bodyStructure_20260901", "en", "clinical/models")
    .setInputCols(Array("bge_embeddings"))
    .setOutputCol("snomed_code")
    .setDistanceFunction("EUCLIDEAN")
    .setThreshold(1000)

val pipeline = new Pipeline().setStages(Array(
    documentAssembler, sentenceDetectorDL, tokenizer, word_embeddings, ner_jsl, ner_jsl_converter, ner_anatomy, ner_anatomy_converter, ner_oncology_anatomy, ner_oncology_anatomy_converter, chunk_merger, chunk2doc, embedder, resolver
))

val data = Seq("The patient is a 30-year-old female with coronary artery disease and swelling affecting the kidney and lower limb.").toDF("text")
val res = pipeline.fit(data).transform(data)

Results

| chunk           | label    |   snomed_code | resolution      | all_codes                                                                           | all_resolutions                                                                     |
|:----------------|:---------|--------------:|:----------------|:------------------------------------------------------------------------------------|:------------------------------------------------------------------------------------|
| coronary artery | BodyPart |      41801008 | coronary artery | 41801008:::119204004:::360487004:::110554000:::51114001:::50018008:::13647002:::... | coronary artery:::coronary artery part:::segment of coronary artery:::coronary a... |
| kidney          | BodyPart |      64033007 | kidney          | 64033007:::119219003:::17373004:::50403003:::58471003:::70948008:::113343008:::9... | kidney:::kidney part:::both kidneys:::cortex of kidney:::renal tubule:::transpla... |
| lower limb      | BodyPart |      61685007 | lower limb      | 61685007:::30021000:::120575009:::128263001:::69548008:::48979004:::244014007:::... | lower limb:::lower leg:::lower limb part:::lower body:::lower body part:::left l... |

Model Information

Model Name: bgeresolve_snomed_bodyStructure_20260901
Compatibility: Healthcare NLP 6.4.1+
License: Licensed
Edition: Official
Input Labels: [bge_embeddings]
Output Labels: [snomed_code]
Language: en
Size: 199.9 MB
Case sensitive: false