Update app.py

app.py

@@ -13,6 +13,36 @@ def reset_conversation():
     st.session_state.messages = []
     return None
 
+def analyze_diagnosis(messages, model_link):
+    # Extract user messages
+    user_messages = [message[1] for message in messages if message[0] == "user"]
+
+    # Define mental conditions and associated keywords
+    mental_conditions = {
+        "Depression": ["depression", "sad", "hopeless", "lonely", "empty", "worthless", "miserable"],
+        "Anxiety": ["anxiety", "nervous", "worried", "fearful", "panicked", "stressed", "tense"],
+        "Panic disorder": ["panic attack", "panic", "scared", "terrified", "frightened", "hyperventilate", "heart racing"],
+        "Bipolar disorder": ["bipolar", "manic", "mania", "euphoric", "energetic", "depressed", "hopeless"],
+        "Schizophrenia": ["schizophrenia", "hallucination", "delusion", "paranoia", "disorganized", "psychotic", "dissociation"],
+        "PTSD": ["ptsd", "trauma", "nightmare", "flashback", "startled", "avoidance", "hypervigilance"],
+        "Obsessive-Compulsive Disorder": ["ocd", "obsession", "compulsion", "intrusive thought", "ritual", "cleaning", "checking"],
+        "Eating disorder": ["eating disorder", "anorexia", "bulimia", "binge eating", "body image", "weight obsession", "purging"],
+        "Substance use disorder": ["substance use", "drug addiction", "alcoholism", "substance abuse", "withdrawal", "craving", "dependency"],
+        "Attention-deficit/hyperactivity disorder (ADHD)": ["adhd", "attention deficit", "hyperactivity", "impulsivity", "inattention", "restlessness", "fidgeting"],
+        "Borderline personality disorder": ["borderline personality", "emotional instability", "impulsivity", "fear of abandonment", "unstable relationships", "self-harm", "mood swings"],
+        "Postpartum depression": ["postpartum depression", "baby blues", "motherhood depression", "crying spells", "irritability", "hopelessness", "fatigue"],
+        "Social anxiety disorder": ["social anxiety", "fear of social situations", "embarrassment", "humiliation", "avoidance", "self-consciousness", "panic"]
+    }
+
+    # Count occurrences of keywords for each mental condition
+    condition_counts = {condition: sum(message.lower().count(keyword) for message in user_messages for keyword in keywords) 
+                        for condition, keywords in mental_conditions.items()}
+
+    # Determine the condition with the highest count
+    diagnosis = max(condition_counts, key=condition_counts.get)
+
+    return diagnosis
+
 # Define model links for the Addiction Recovery and Mental Health models
 model_links = {
     "Addiction recovery AI": "NousResearch/Nous-Hermes-2-Yi-34B",
@@ -123,4 +153,10 @@ if prompt := st.chat_input(f"Hi, I'm {selected_model}, let's chat"):
                         st.write("You can reach out to a certified therapist at +25493609747.")
         
         # Add assistant response to chat history
-        st.session_state.messages.append(("assistant", assistant_response))
+        st.session_state.messages.append(("assistant", assistant_response))
+
+# Add diagnostic feature to the sidebar
+diagnosis = analyze_diagnosis(st.session_state.messages, model_links[selected_model])
+if diagnosis:
+    st.sidebar.subheader("Diagnosis")
+    st.sidebar.write(diagnosis)