Update app.py

app.py

@@ -1,18 +1,14 @@
 import streamlit as st
-from transformers import pipeline, GPT2LMHeadModel, GPT2Tokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 from gtts import gTTS
-from pytube import Search
-import os
 import random
+from youtubesearchpython import Search
 
-# Initialize GPT-2 model and tokenizer from Hugging Face
-tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
-model = GPT2LMHeadModel.from_pretrained("gpt2")
+# Load DialoGPT model and tokenizer from Hugging Face
+tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
+model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
 
-# Create a text generation pipeline using GPT-2
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
-
-# Set up Streamlit page
+# Set up Streamlit page configuration
 st.set_page_config(page_title="Grief and Loss Support Bot", page_icon="🌿", layout="centered")
 st.markdown("""
     <style>
@@ -25,61 +21,38 @@ st.markdown("""
     </style>
 """, unsafe_allow_html=True)
 
-# Title
+# Title and introduction to the bot
 st.title("Grief and Loss Support Bot 🌿")
 st.subheader("Your compassionate companion in tough times 💚")
 
-# Get user input
+# User input
 user_input = st.text_input("Share what's on your mind...", placeholder="Type here...", max_chars=500)
 
 # Store previous responses to check for repetition
 if 'previous_responses' not in st.session_state:
     st.session_state.previous_responses = []
 
-# Check if user has entered text
-if user_input:
-    # Refined prompt for better empathetic responses
-    prompt = f"User is feeling overwhelmed with emotional distress and is going through a tough time. Respond empathetically, offering support and understanding. The user's input: {user_input}"
-
-    # Run the text generation model to generate a response based on user input
-    generated_responses = generator(prompt, max_length=200, num_return_sequences=3, temperature=0.7)
-
-    # Filter out any responses that are too similar to previous responses or user input
-    new_responses = [response['generated_text'].strip() for response in generated_responses]
-    new_responses = [resp for resp in new_responses if resp.lower() not in [prev.lower() for prev in st.session_state.previous_responses] and resp.lower() != user_input.lower()]
-
-    # If there are valid new responses, pick one, otherwise fallback
-    if new_responses:
-        selected_response = random.choice(new_responses)
-    else:
-        # If no new response, fallback to a more generic empathetic message
-        fallback_responses = [
-            "I understand how you're feeling. You're not alone in this. I'm here to listen and help.",
-            "I'm really sorry you're going through this. Let's take one step at a time. I'm here for you.",
-            "It sounds really tough right now. It's okay to feel overwhelmed. You're doing your best, and that's enough."
-        ]
-        selected_response = random.choice(fallback_responses)
-
-    # Add extra empathetic phrases to the response
-    extra_empathy = [
-        "It’s completely normal to feel this way when things get tough. You're doing great by reaching out.",
-        "I know it can feel like a lot right now, but one step at a time. You're not alone in this.",
-        "Even in the toughest times, remember that there’s always support around you."
-    ]
-    selected_response += " " + random.choice(extra_empathy)
-
-    # Store the new response for future checks
-    st.session_state.previous_responses.append(selected_response)
-
-    # Display response
-    st.text_area("Bot's Response:", selected_response, height=250)
-
-    # Text-to-speech output
-    tts = gTTS(selected_response, lang='en')
-    audio_file = "response.mp3"
-    tts.save(audio_file)
-    st.audio(audio_file, format="audio/mp3")
-
+# Function to generate a more empathetic and focused response using DialoGPT
+def generate_response(user_input):
+    # Encode the input text and generate a response
+    new_user_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors='pt')
+    bot_input_ids = new_user_input_ids
+    chat_history_ids = model.generate(bot_input_ids, max_length=200, pad_token_id=tokenizer.eos_token_id, temperature=0.7, top_k=50, repetition_penalty=1.2)
+    
+    # Decode the response to text
+    chat_history_ids = chat_history_ids[:, bot_input_ids.shape[-1]:]  # remove the input from the response
+    bot_output = tokenizer.decode(chat_history_ids[0], skip_special_tokens=True)
+    
+    # Build a more empathetic and thoughtful response
+    response = f"I’m really sorry you're feeling like this. It’s completely normal to feel overwhelmed when you're facing a heavy workload. It’s important to acknowledge how you feel and not keep it bottled up. Sometimes, stress and emotional exhaustion can build up, and it’s okay to let yourself feel those emotions."
+
+    # Add coping strategies based on the situation
+    if "workload" in user_input.lower():
+        response += "\n\nWhen the workload feels too heavy, it can be helpful to break tasks down into smaller, more manageable steps. Focus on one thing at a time, and remember that it’s okay to take breaks when needed. Asking for support from colleagues or friends is also a good way to lighten the load."
+
+    # Add general supportive message
+    response += "\n\nYou're doing your best, and that’s all anyone can ask for. Please take care of yourself and know that it’s okay to take a step back when things feel too much. Your well-being is the most important thing."
+    
     # Suggest a productive activity based on detected keywords
     if any(keyword in user_input.lower() for keyword in ["lonely", "lost", "sad", "overwhelmed", "academic", "exam"]):
         st.info("Here's a suggestion to help you cope:")
@@ -111,3 +84,20 @@ if user_input:
     if any(keyword in user_input.lower() for keyword in crisis_keywords):
         st.warning("It seems like you might be in distress. Please reach out to a crisis hotline or a trusted individual.")
         st.write("[Find emergency resources here](https://www.helpguide.org/find-help.htm)")
+
+    return response
+
+# Check if the user has typed something
+if user_input:
+    # Generate the empathetic response
+    response = generate_response(user_input)
+    
+    # Store and show the new response
+    st.session_state.previous_responses.append(response)
+    st.text_area("Bot's Response:", response, height=250)
+
+    # Text-to-speech output (optional)
+    tts = gTTS(response, lang='en')
+    audio_file = "response.mp3"
+    tts.save(audio_file)
+    st.audio(audio_file, format="audio/mp3")