Update app.py

app.py

@@ -2,63 +2,30 @@ import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from gtts import gTTS
 from pytube import Search
-import random
 import os
 
 # Load pretrained models
 tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
 model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
-# Load Llama 2 model
-llama_model_name = "meta-llama/Llama-2-7b-hf"
-llama_tokenizer = AutoTokenizer.from_pretrained(llama_model_name)
-llama_model = AutoModelForCausalLM.from_pretrained(llama_model_name)
+# Load GPT-2 model and tokenizer for story generation
+gpt2_tokenizer = AutoTokenizer.from_pretrained("gpt2-medium")
+gpt2_model = AutoModelForCausalLM.from_pretrained("gpt2-medium")
+
 emotion_classifier = pipeline("text-classification", model="bhadresh-savani/distilbert-base-uncased-emotion", return_all_scores=True)
 
+# Function to generate a comforting story using GPT-2
 def generate_story(theme):
-    # A detailed prompt for generating a comforting story about the selected theme
-    story_prompt = f"Write a comforting, detailed, and heartwarming story about {theme}. The story should include a character who faces a tough challenge, finds hope, and ultimately overcomes the situation with a positive resolution."
-    
-    # Generate story using Llama 2
-    input_ids = llama_tokenizer.encode(story_prompt, return_tensors='pt')
-    
-    story_ids = llama_model.generate(
-        input_ids,
-        max_length=600,  # Generate longer stories
-        temperature=0.7,  # Balanced creativity
-        top_p=0.9,
-        repetition_penalty=1.1,
-        num_return_sequences=1
-    )
-    
-    # Decode the generated text
-    story = llama_tokenizer.decode(story_ids[0], skip_special_tokens=True)
-    return story
-
-
+    story_prompt = f"Write a comforting, detailed, and heartwarming story about {theme}. Include a character facing challenges and finding hope."
+    input_ids = gpt2_tokenizer.encode(story_prompt, return_tensors='pt')
+    story_ids = gpt2_model.generate(input_ids, max_length=500, temperature=0.8, top_p=0.9, repetition_penalty=1.2)
+    return gpt2_tokenizer.decode(story_ids[0], skip_special_tokens=True)
 
 # Function to generate an empathetic response
-# Function to generate an empathetic response using Llama 2
 def generate_response(user_input):
-    # Create a prompt for the Llama 2 model to generate an empathetic response
-    response_prompt = f"You are a compassionate support bot. A user has shared: '{user_input}'. Rather than sympathy, respond with empathy, understanding, and encouragement in a supportive manner."
-    
-    # Tokenize the prompt
-    input_ids = llama_tokenizer.encode(response_prompt, return_tensors='pt')
-    
-    # Generate the response using Llama 2
-    response_ids = llama_model.generate(
-        input_ids,
-        max_length=300,        # Generate a longer response
-        temperature=0.7,        # Adjust temperature for coherent and empathetic responses
-        top_p=0.9,              # Use top-p sampling for balanced creativity
-        repetition_penalty=1.1, # Avoid repetitive responses
-        num_return_sequences=1
-    )
-    
-    # Decode the generated text
-    response = llama_tokenizer.decode(response_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
-    return response
-
+    response_prompt = f"You are a compassionate support bot. A user has shared: '{user_input}'. Respond with empathy and encouragement."
+    input_ids = tokenizer.encode(response_prompt, return_tensors='pt')
+    chat_history_ids = model.generate(input_ids, max_length=300, temperature=0.85, top_k=50, repetition_penalty=1.2)
+    return tokenizer.decode(chat_history_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
 
 # Analyze user input for emotional tone
 def get_emotion(user_input):
@@ -68,13 +35,12 @@ def get_emotion(user_input):
 
 # Function to fetch YouTube videos
 def fetch_youtube_videos(activity):
-    search = Search(f"{activity} for mental health relaxation")
-    search_results = search.results[:3]
-    videos = []
-    for video in search_results:
-        video_url = f"https://www.youtube.com/watch?v={video.video_id}"
-        videos.append((video.title, video_url))
-    return videos
+    try:
+        search = Search(f"{activity} for mental health relaxation")
+        videos = [(video.title, f"https://www.youtube.com/watch?v={video.video_id}") for video in search.results[:3]]
+        return videos
+    except Exception:
+        return []
 
 # Streamlit page configuration
 st.set_page_config(page_title="Grief and Loss Support Bot 🌿", page_icon="🌿", layout="centered")
@@ -93,7 +59,6 @@ with st.sidebar:
             tts.save(meditation_audio)
         st.audio(meditation_audio, format="audio/mp3")
 
-# Generating a comforting story
 st.sidebar.header("📖 Short Comforting Story")
 story_theme = st.selectbox("Choose a theme for your story:", ["courage", "healing", "hope"])
 if st.sidebar.button("Generate Story"):
@@ -101,35 +66,29 @@ if st.sidebar.button("Generate Story"):
         story = generate_story(story_theme)
     st.text_area("Here's your story:", story, height=300)
 
-
-
 # User input section
 user_input = st.text_input("Share what's on your mind...", placeholder="Type here...", max_chars=500)
 
 # Initialize session state
-if 'previous_responses' not in st.session_state:
-    st.session_state.previous_responses = []
 if 'badges' not in st.session_state:
     st.session_state.badges = []
 
 if user_input:
     with st.spinner("Thinking..."):
+        emotion = get_emotion(user_input)
         response = generate_response(user_input)
+        st.text_area("Bot's Response:", response, height=250)
     
-    # Display the bot's response
-    st.text_area("Bot's Response:", response, height=250)
-
-    
-    # Assign motivational badges
+    # Assign badges based on emotional analysis
     if emotion in ["joy", "optimism"]:
         badge = "🌟 Positivity Badge"
         if badge not in st.session_state.badges:
             st.session_state.badges.append(badge)
             st.success(f"Congratulations! You've earned a {badge}!")
-    
+
     # Suggest activities
     st.info("🎨 Try a New Activity")
-    activities = ["exercise", "yoga", "journaling", "painting", "meditation"]
+    activities = ["exercise", "yoga", "journaling", "painting", "meditation", "swimming"]
     selected_activity = st.selectbox("Pick an activity:", activities)
 
     if st.button("Find Videos"):
@@ -145,14 +104,14 @@ if any(word in user_input.lower() for word in ["suicide", "help", "depressed"]):
     st.warning("Please reach out to a crisis hotline for immediate support.")
     st.write("[Find emergency resources here](https://www.helpguide.org/find-help.htm)")
 
-# Generate audio response
+# Generate and play audio response
 if user_input:
     tts = gTTS(response, lang='en')
     audio_file = "response.mp3"
     tts.save(audio_file)
     st.audio(audio_file, format="audio/mp3")
 
-# Display badges
+# Display earned badges
 if st.session_state.badges:
     st.sidebar.header("🏅 Achievements")
     for badge in st.session_state.badges: