ZeroShot Multitask Oncology

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

zero_shot = PretrainedZeroShotMultiTask.pretrained("zeroshot_multitask_oncology", "en", "clinical/models")\
    .setInputCols(["document"])\
    .setOutputCol("extractions")\
    .setEntityThreshold(0.4)\
    .setEntities([
        "Adenopathy::Pathological findings of the lymph nodes",
        "Age::All mentions of ages related to the patient or anyone else",
        "Biomarker::Biological molecules indicating presence, absence, or type of cancer (oncogenes excluded)",
        "Biomarker_Result::Terms or values identified as the result of a biomarker",
        "Cancer_Dx::Cancer diagnoses or pathological types used as synonyms for cancer (e.g. carcinoma); anatomical references included",
        "Cancer_Score::Clinical or imaging scores specific to cancer settings (e.g. BI-RADS, Allred score)",
        "Cancer_Surgery::Terms indicating surgery as a form of cancer treatment",
        "Chemotherapy::Chemotherapy drugs or unspecific words such as chemotherapy",
        "Cycle_Count::Total number of cycles of an oncological therapy (e.g. 5 cycles)",
        "Cycle_Day::References to the day of the cycle of oncological therapy (e.g. day 5)",
        "Cycle_Number::The number of the cycle of an oncological therapy being applied (e.g. third cycle)",
        "Date::Exact dates in any format, including day number, month and/or year",
        "Death_Entity::Words indicating death of the patient or someone else (e.g. died, passed away)",
        "Direction::Directional and laterality terms such as left, right, bilateral, upper, lower",
        "Dosage::Quantity prescribed by the physician for an active ingredient",
        "Duration::Words indicating the duration of a treatment (e.g. for 2 weeks)",
        "Frequency::Words indicating the frequency of treatment administration (e.g. daily, bid)",
        "Gender::Gender-specific nouns and pronouns, including family members such as father",
        "Grade::Pathological grading of tumors or degrees of cellular differentiation",
        "Histological_Type::Histological variants or cancer subtypes (e.g. papillary, clear cell, medullary)",
        "Hormonal_Therapy::Hormonal drugs used to treat cancer, or unspecific words such as hormonal therapy",
        "Imaging_Test::Imaging tests mentioned in texts (e.g. chest CT scan)",
        "Immunotherapy::Immunotherapy drugs or unspecific words such as immunotherapy",
        "Invasion::Mentions referring to tumor invasion such as invasion or involvement (metastases excluded)",
        "Line_Of_Therapy::Explicit references to the line of therapy (e.g. first-line treatment)",
        "Metastasis::Terms indicating metastatic disease (anatomical references not included)",
        "Oncogene::Mentions of genes implicated in the etiology of cancer",
        "Pathology_Result::Biopsy findings from pathology report not covered by another entity (e.g. malignant ductal cells)",
        "Pathology_Test::Mentions of biopsies or tests that use tissue samples",
        "Performance_Status::Performance status scores such as ECOG and Karnofsky, including the result (e.g. ECOG performance status of 4)",
        "Race_Ethnicity::Racial, national origin, or sociocultural group categories",
        "Radiotherapy::Terms indicating the use of radiotherapy",
        "Response_To_Treatment::Terms related to clinical progress of the patient related to cancer treatment (e.g. recurrence, bad response, improvement)",
        "Relative_Date::Temporal references relative to the date of the text or another specific date (e.g. yesterday, three years later)",
        "Route::Words indicating the type of administration route (e.g. PO, transdermal)",
        "Site_Bone::Anatomical terms referring to the human skeleton",
        "Site_Brain::Anatomical terms referring to the central nervous system including brain stem and cerebellum",
        "Site_Breast::Anatomical terms referring to the breasts",
        "Site_Liver::Anatomical terms referring to the liver",
        "Site_Lung::Anatomical terms referring to the lungs",
        "Site_Lymph_Node::Anatomical terms referring to lymph nodes, excluding adenopathies",
        "Site_Other_Body_Part::Relevant anatomical terms not included in other anatomical entities",
        "Smoking_Status::All mentions of smoking related to the patient or someone else",
        "Staging::Mentions of cancer stage (e.g. stage 2b, T2N1M0, in situ, early-stage, advanced)",
        "Targeted_Therapy::Targeted therapy drugs or unspecific words such as targeted therapy",
        "Tumor_Finding::Nonspecific terms related to tumors, malignant or benign (e.g. mass, tumor, lesion, neoplasm)",
        "Tumor_Size::Size of the tumor including numerical value and unit of measurement (e.g. 3 cm)",
        "Unspecific_Therapy::Known cancer therapy not specific to any other therapy entity (e.g. chemoradiotherapy, adjuvant therapy)",
    ])    .setClassifications([
        ("cancer_type",       ["breast", "lung", "colorectal", "prostate", "hematologic", "other", "not_cancer"]),
        ("staging_category",  ["early_stage", "locally_advanced", "metastatic", "unknown"]),
    ])    .setRelations([
        "therapies_treats_cancer_dx",
        "therapies_treats_tumor_finding",
        "therapies_Response_To_Treatmentt",
        "imaging_test_reveals_finding",
        "biomarker_associated_with_cancer_dx_or_tumor_finding",
        "oncogene_associated_with_cancer_dx",
    ])

pipeline = Pipeline(
    stages = [
        document_assembler,
        zero_shot
])

text = f"""A 58-year-old female was diagnosed with stage IIIA HER2-positive invasive ductal carcinoma of the right breast. She received 6 cycles of neoadjuvant carboplatin plus docetaxel followed by mastectomy. Adjuvant trastuzumab 6 mg/kg IV every 3 weeks was initiated for 12 months."""

data = spark.createDataFrame([[text]]).toDF("text")

results = pipeline.fit(data).transform(data)

results.select("extractions").show(truncate=False)