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

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 = MPNetEmbeddings
    .pretrained("mpnet_embeddings_biolord_2023_c", "en")
    .setInputCols(Array("ner_chunk_doc"))
    .setOutputCol("resolver_embeddings")
    .setCaseSensitive(false)
    .setBatchSize(1)

val resolver = SentenceEntityResolverModel
    .pretrained("biolordresolve_snomed_bodyStructure_20260901", "en", "clinical/models")
    .setInputCols(Array("resolver_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)