Smoking Usage For Classification

document_assembler = DocumentAssembler()\
    .setInputCol("text")\
    .setOutputCol("document")
        
sent_embd = MPNetEmbeddings.pretrained("mpnet_embedding_all_mpnet_base_v2_by_sentence_transformers", 'en')\
    .setInputCols(["document"])\
    .setOutputCol("sentence_embeddings")\

features_asm = FeaturesAssembler()\
    .setInputCols(["sentence_embeddings"])\
    .setOutputCol("features")
      
generic_classifier = GenericClassifierModel.pretrained('genericclassifier_smoking_mpnet_wip', 'en', 'clinical/models')\
    .setInputCols("features")\
    .setOutputCol("prediction")\

pipeline = Pipeline(stages=[
        document_assembler,
        sent_embd,
        features_asm,
        generic_classifier])

text_list = [
    """Mr. John Smith, a 45-year-old truck driver, visited the clinic with enduring symptoms of cough and breathlessness, which have progressively worsened over recent months. He denies any significant past medical history, although he is a current smoker, consuming ten cigarettes daily for the last two decades.""",  
    """History of Present Illness: 52-year-old teacher Mrs. Sarah Johnson arrived to the clinic complaining of a chronic cough and dyspnea that had been becoming worse over the previous few months. She is a current smoker, and has smoked for the past 30 years, averaging a pack of cigarettes every day.""",
    """Pulmonary Function Tests: showed restricted airflow and decreased FEV1 and FEV1/FVC ratios, which are indicators of chronic obstructive pulmonary disease (COPD). Diagnosis: Smoking-related chronic obstructive pulmonary disease (COPD)""",
    """Ms. Jane Doe, a 60-year-old retired teacher, presented to the emergency department complaining of severe abdominal pain and vomiting. She has a history of gallstones but has been asymptomatic for years. Currently, she does not smoke or drink alcohol, focusing on a healthy lifestyle.""",
    """On admission, he reported feeling anxious, agitated, and excessive sweating but denied nausea, vomiting, headache, auditory, visual, or tactile hallucinations. The system review was otherwise unfavorable. He denied smoking and using illegal drugs, and he was not taking any medication.""",
]

df = spark.createDataFrame(text_list, StringType()).toDF("text")

result = pipeline.fit(df).transform(df)

result.select("text", "prediction.result").show(truncate=100)