Detect Assertion Status from Drug Entities

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

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

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

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

val ner_model = MedicalNerModel.pretrained("ner_opioid", "en", "clinical/models")
    .setInputCols(Array("sentence", "token", "embeddings"))
    .setOutputCol("ner")

val ner_converter = new NerConverterInternal()
    .setInputCols(Array("sentence", "token", "ner"))
    .setOutputCol("ner_chunk")
    .setWhiteList(Array("opioid_drug", "other_drug"))


val assertion = AssertionDLModel.pretrained("assertion_opioid_drug_status_wip", "en", "clinical/models")
    .setInputCols(Array("sentence", "ner_chunk", "embeddings"))
    .setOutputCol("assertion")
    
val pipeline = new Pipeline().setStages(Array(document_assembler,
                                              sentence_detector,
                                              tokenizer,
                                              word_embeddings,
                                              ner_model,
                                              ner_converter,
                                              assertion))

val data = Seq("""The patient presented to the hospital for a neurological evaluation, with a documented prescription for Percocet to manage chronic back pain. Assessment revealed ongoing discomfort localized to the lumbar region, with associated numbness and tingling in the lower extremities.""",
               """The patient, with a known history of hypertension managed with atenolol 50mg and verapamil 40mg, presented after a fall resulting in an ankle injury. Examination revealed swelling and tenderness, indicative of a twisted ankle. Considering the patient's medical history and pain management needs, a prescription for tramadol was provided to alleviate discomfort while ensuring minimal impact on blood pressure control.""",
               """The patient presented to the rehabilitation facility with a documented history of opioid abuse, primarily stemming from misuse of prescription percocet pills intended for their partner's use. Initial assessment revealed withdrawal symptoms consistent with opioid dependency, including agitation, diaphoresis, and myalgias.""",
               """The patient presented to the emergency department following an overdose on cocaine. On examination, the patient displayed signs of sympathetic nervous system stimulation, including tachycardia, hypertension, dilated pupils, and agitation.""",
               """The patient, with a documented history of chronic pain syndrome, was admitted following an accidental overdose of prescribed OxyContin. Upon assessment, the patient displayed symptoms indicative of opioid toxicity, including respiratory depression, altered mental status, and pinpoint pupils. Immediate resuscitative measures were undertaken, including airway management, administration of naloxone, and close monitoring of vital signs.""").toDF("text")


val result = pipeline.fit(data).transform(data)