Update app.py

app.py

@@ -1,40 +1,36 @@
-# Use a pipeline as a high-level helper
 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")
 
-
-# Streamlit app with background image
-def st_display_background(image_url):
-    st_style = f"""
-            <style>
-            body {{
-                background-image: url("{image_url}") !important;
-                background-size: cover;
-            }}
-            </style>
-            """
-    st.markdown(st_style, unsafe_allow_html=True)
-
-image_url = "https://images.unsplash.com/photo-1504701954957-2010ec3bcec1?q=80&w=1974&auto=format&fit=crop&ixlib=rb-4.0.3&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D"
-
-
-def ensemble_predict(text): 
+# Function for ensemble prediction
+def ensemble_predict(text):
     # Store results from each model
     results_1 = pipe_1(text)    
     results_2 = pipe_2(text)
 
-    ensemble_scores = {} 
-    
-    for results in [results_1, results_2]:  # Iterate through predictions
+    # 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] = ensemble_scores.get(label, 0) + score / 2  
+            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
 
@@ -46,21 +42,20 @@ 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:")
 
-st_display_background(image_url)
-
 if st.button('Predict'):
-  # ... (input validation ... )
-  predicted_label, confidence = ensemble_predict(sentence)
-
-  # CSS injection to target the labels
-  st.markdown("""
-  <style>
-     div[data-testid="metric-container"]  {
-        font-weight: bold;
-        font-size: 18px; /* Adjust the font size as desired */
-     }
-  </style>
-  """, unsafe_allow_html=True)  
-
-  st.write("Result:", predicted_label)
-  st.write("Confidence:", confidence)
+    # Perform the prediction
+    predicted_label, confidence = ensemble_predict(sentence)
+
+    # CSS injection to target the labels
+    st.markdown("""
+    <style>
+        div[data-testid="metric-container"]  {
+            font-weight: bold;
+            font-size: 18px; /* Adjust the font size as desired */
+        }
+    </style>
+    """, unsafe_allow_html=True)  
+
+    # Display the result
+    st.write("Result:", predicted_label)
+    st.write("Confidence:", confidence)