ZeroShot Multitask Oncology

from johnsnowlabs import nlp, medical
document_assembler = nlp.DocumentAssembler()\
    .setInputCol("text")\
    .setOutputCol("document")

zero_shot = medical.PretrainedZeroShotMultiTask.pretrained("zeroshot_multitask_oncology_generic", "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_Result::Terms or values identified as the result of a biomarker",
        "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",
        "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)",
        "Grade::Pathological grading of tumors or degrees of cellular differentiation",
        "Histological_Type::Histological variants or cancer subtypes (e.g. papillary, clear cell, medullary)",
        "Invasion::Mentions referring to tumor invasion such as invasion or involvement (metastases excluded)",
        "Metastasis::Terms indicating metastatic disease (anatomical references not included)",
        "Pathology_Result::Biopsy findings from pathology report not covered by another entity (e.g. malignant ductal cells)",
        "Performance_Status::Performance status scores such as ECOG and Karnofsky, including the result (e.g. ECOG performance status of 4)",
        "Radiation_Dose::Radiation dose values used in radiotherapy (e.g. 40 Gy, 3000 cGy)",
        "Response_To_Treatment::Terms related to clinical progress of the patient related to cancer treatment (e.g. recurrence, bad response, improvement)",
        "Route::Words indicating the type of administration route (e.g. PO, transdermal)",
        "Staging::Mentions of cancer stage (e.g. stage 2b, T2N1M0, in situ, early-stage, advanced)",
        "Tumor_Size::Size of the tumor including numerical value and unit of measurement (e.g. 3 cm)",
        "Biomarker::Biological molecules and genes indicating presence, absence, or type of cancer, including oncogenes (e.g. EGFR, HER2, BRCA1)",
        "Body_Part::Anatomical locations and body parts including bones, brain, breast, liver, lungs, lymph nodes, and other body regions",
        "Cancer_Tumor::Cancer diagnoses, pathological cancer types, and nonspecific tumor findings (e.g. carcinoma, adenocarcinoma, mass, lesion, neoplasm)",
        "Cycle_Info::References to oncological therapy cycles including total cycle count (e.g. 5 cycles), cycle day (e.g. day 5), and cycle number (e.g. third cycle)",
        "Date::Date and temporal expressions including exact dates (e.g. March 2023) and relative references (e.g. yesterday, three years later)",
        "Demographics::Patient demographic information including gender-specific terms, racial or ethnic background, and smoking history",
        "Test::Diagnostic procedures including imaging studies (e.g. CT scan, MRI, PET scan) and pathology tests using tissue samples (e.g. biopsy)",
        "Therapy::Any cancer treatment including chemotherapy drugs, hormonal therapy, immunotherapy, targeted therapy, radiotherapy, and unspecific therapies; also includes line of therapy references (e.g. first-line treatment)",
    ])\
    .setClassifications([
        ("cancer_type",      ["breast", "lung", "colorectal", "prostate", "hematologic", "other", "not_cancer"]),
        ("staging_category", ["early_stage", "locally_advanced", "metastatic", "unknown"]),
    ])\
    .setRelations([
        "therapy_treats_cancer_tumor",
        "therapy_causes_adverse_effect",
        "test_reveals_finding",
        "biomarker_associated_with_cancer",
        "biomarker_predicts_response_to_therapy",
    ])

pipeline = nlp.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)