Text Classification for Age

document_assembler = DocumentAssembler()\
    .setInputCol("text")\
    .setOutputCol("document")
        
embeddings = E5Embeddings.pretrained("e5_large", "en")\
    .setInputCols(["document"])\
    .setOutputCol("embeddings")\
    .setMaxSentenceLength(512)

features_asm = FeaturesAssembler()\
    .setInputCols(["embeddings"])\
    .setOutputCol("features")

generic_classifier = GenericClassifierModel.pretrained("genericclassifier_age_e5", "en", "clinical/models")\
    .setInputCols(["features"])\
    .setOutputCol("prediction")

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

sample_texts = [
"""The patient presents with conditions often associated with the stresses and lifestyle of early career and possibly higher education stages, including sleep irregularities and repetitive stress injuries. There's a notable emphasis on preventative care, with discussions around lifestyle choices that can impact long-term health, such as smoking cessation, regular exercise, and balanced nutrition. The patient is also counseled on mental health, particularly in managing stress and anxiety that may arise from personal and professional responsibilities and ambitions at this stage of life.""",
"""The patient, in the infancy stage, is experiencing typical developmental milestones such as teething and learning to crawl. Regular check-ups focus on monitoring growth patterns, ensuring proper nutrition, especially with the transition from milk to solid foods, and addressing common infant concerns like colic and sleep patterns. Immunizations are up-to-date, following the recommended schedule for this critical early life stage.""",
"""The toddler patient shows normal signs of curiosity and motor skill development for their age. Concerns include potential delays in speech development and the challenges of transitioning from crib to bed. Nutritional guidance is provided, focusing on balanced diets and managing the picky eating phase common in toddlers. Safety in the home environment is also discussed, given the patient's increased mobility and exploratory behavior.""",
"""The senior patient presents with age-related issues such as reduced hearing and vision, arthritis, and memory lapses. Emphasis is on managing chronic conditions, maintaining social engagement, and adapting lifestyle to changing physical abilities. Discussions include medication management, dietary adjustments to suit older digestion, and the importance of regular, low-impact exercise.""",
"""The late teenage patient is dealing with final growth spurts, the stress of impending adulthood, and decisions about higher education or career paths. Health discussions include maintaining a balanced diet, the importance of regular sleep patterns, and managing academic and social pressures. Mental health support is considered crucial at this stage, with a focus on building resilience and coping mechanisms.""",
"""The patient, faces adjustments to a new lifestyle with changes in daily routines and social interactions. Health concerns include managing the transition from an active work life to more leisure time, which may impact physical and mental health. Preventative health measures are emphasized, along with the importance of staying mentally and physically active and engaged in the community."""
]

data = spark.createDataFrame(sample_texts, StringType()).toDF("text")

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