from transformers import pipeline
import streamlit as st
import streamlit.components.v1 as components

# Load the models
pipe_1 = pipeline("text-classification", model="mavinsao/roberta-base-finetuned-mental-health")
pipe_2 = pipeline("text-classification", model="mavinsao/mi-roberta-base-finetuned-mental-health")

# Function for ensemble prediction
def ensemble_predict(text):
    # Store results from each model
    results_1 = pipe_1(text)    
    results_2 = pipe_2(text)

    # Initialize a dictionary with all potential labels to ensure they are considered
    ensemble_scores = {}

    # Add all labels from the first model's output
    for result in results_1:
        ensemble_scores[result['label']] = 0

    # Add all labels from the second model's output
    for result in results_2:
        ensemble_scores[result['label']] = 0

    # Aggregate scores from both models
    for results in [results_1, results_2]:  
        for result in results:
            label = result['label']
            score = result['score']
            ensemble_scores[label] += score / 2  # Averaging the scores

    # Determine the predicted label and confidence
    predicted_label = max(ensemble_scores, key=ensemble_scores.get) 
    confidence = ensemble_scores[predicted_label]  # Ensemble confidence

    return predicted_label, confidence

# Streamlit app
st.title('Mental Illness Prediction')

# Input text area for user input
sentence = st.text_area("Enter the long sentence to predict your mental illness state:")

if st.button('Predict'):
    # Perform the prediction
    predicted_label, confidence = ensemble_predict(sentence)

 # Display the result with custom styling
    st.markdown(f"""
    <h2 style='text-align: center; color: #1E90FF;'>Prediction Results</h2>
    <p style='font-size: 24px; font-weight: bold;'>Result: <span style='color: #1E90FF;'>{predicted_label}</span></p>
    <p style='font-size: 24px; font-weight: bold;'>Confidence: <span style='color: #1E90FF;'>{confidence:.2f}</span></p>
    """, unsafe_allow_html=True)

st.info("Remember: This prediction is not a diagnosis. Our method is designed to support, not replace, mental health professionals. The model's predictions should be used as a reference, and the final diagnosis should be made by a qualified professional to avoid potential biases and inaccuracies.")