Mapping Entities (Disease or Syndrome) with Corresponding UMLS CUI Codes

Description

This model maps disease and syndrome entities extracted by NER to UMLS CUI codes covering the Disease or Syndrome semantic type (T047), comprising approximately 399,000 name–CUI pairs. It is trained on the 2026AA release of the Unified Medical Language System (UMLS) dataset.

Live Demo Open in Colab Copy S3 URI

How to use


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

sentence_detector = SentenceDetector()\
    .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_model = MedicalNerModel.pretrained("ner_clinical_large", "en", "clinical/models")\
    .setInputCols(["sentence", "token", "embeddings"])\
    .setOutputCol("clinical_ner")

ner_model_converter = NerConverterInternal()\
    .setInputCols(["sentence", "token", "clinical_ner"])\
    .setOutputCol("ner_chunk")

chunkerMapper = ChunkMapperModel.pretrained("umls_disease_syndrome_mapper", "en", "clinical/models")\
    .setInputCols(["ner_chunk"])\
    .setOutputCol("mappings")\
    .setRels(["umls_code"])\
    .setLowerCase(True)

mapper_pipeline = Pipeline(stages=[
    document_assembler,
    sentence_detector,
    tokenizer,
    word_embeddings,
    ner_model,
    ner_model_converter,
    chunkerMapper,
])
data = spark.createDataFrame([["A 35-year-old male with a history of obesity and gestational diabetes mellitus and acyclovir allergy. The patient was diagnosed with Parkinson disease and sepsis."]]).toDF("text")
result = mapper_pipeline.fit(data).transform(data)


document_assembler = nlp.DocumentAssembler()\
    .setInputCol("text")\
    .setOutputCol("document")

sentence_detector = nlp.SentenceDetector()\
    .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_model = medical.NerModel.pretrained("ner_clinical_large", "en", "clinical/models")\
    .setInputCols(["sentence", "token", "embeddings"])\
    .setOutputCol("clinical_ner")

ner_model_converter = medical.NerConverterInternal()\
    .setInputCols(["sentence", "token", "clinical_ner"])\
    .setOutputCol("ner_chunk")

chunkerMapper = medical.ChunkMapperModel.pretrained("umls_disease_syndrome_mapper", "en", "clinical/models")\
    .setInputCols(["ner_chunk"])\
    .setOutputCol("mappings")\
    .setRels(["umls_code"])\
    .setLowerCase(True)

mapper_pipeline = nlp.Pipeline(stages=[
    document_assembler,
    sentence_detector,
    tokenizer,
    word_embeddings,
    ner_model,
    ner_model_converter,
    chunkerMapper,
])
data = spark.createDataFrame([["A 35-year-old male with a history of obesity and gestational diabetes mellitus and acyclovir allergy. The patient was diagnosed with Parkinson disease and sepsis."]]).toDF("text")
result = mapper_pipeline.fit(data).transform(data)


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

val sentenceDetector = new SentenceDetector()
    .setInputCols(Array("document"))
    .setOutputCol("sentence")

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

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

val nerModel = MedicalNerModel.pretrained("ner_clinical_large", "en", "clinical/models")
    .setInputCols(Array("sentence", "token", "embeddings"))
    .setOutputCol("clinical_ner")

val nerModelConverter = new NerConverterInternal()
    .setInputCols(Array("sentence", "token", "clinical_ner"))
    .setOutputCol("ner_chunk")

val chunkerMapper = ChunkMapperModel.pretrained("umls_disease_syndrome_mapper", "en", "clinical/models")
    .setInputCols(Array("ner_chunk"))
    .setOutputCol("mappings")
    .setRels(Array("umls_code"))
    .setLowerCase(true)

val mapperPipeline = new Pipeline().setStages(Array(
    documentAssembler,
    sentenceDetector,
    tokenizer,
    wordEmbeddings,
    nerModel,
    nerModelConverter,
    chunkerMapper,
))

val data = Seq("A 35-year-old male with a history of obesity and gestational diabetes mellitus and acyclovir allergy. The patient was diagnosed with Parkinson disease and sepsis.").toDF("text")
val result = mapperPipeline.fit(data).transform(data)

Results

| ner_chunk                     | umls_code   |
|:------------------------------|:------------|
| obesity                       | C0028754    |
| gestational diabetes mellitus | C0085207    |
| acyclovir allergy             | C0571297    |
| Parkinson disease             | C0030567    |
| sepsis                        | C0036690    |

Model Information

Model Name: umls_disease_syndrome_mapper
Compatibility: Healthcare NLP 6.4.0+
License: Licensed
Edition: Official
Input Labels: [ner_chunk]
Output Labels: [mappings]
Language: en
Size: 12.6 MB