Update app.py

app.py

@@ -58,7 +58,10 @@ def main():
 
         # Make prediction
         response = model.predict(processed_text)
-        predicted_class = response.argmax(axis=-1) 
+        predicted_class = response.argmax(axis=-1)[0] 
+        # Map the predicted class index to the corresponding label
+        predicted_label = class_labels.get(predicted_class, "Unknown")
+                
         # Display the prediction result
         st.write("Predicted Mental State:", predicted_class)