Update app.py

app.py

@@ -8,22 +8,33 @@ import os
 # Load pretrained models
 tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
 model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
+# 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 the pretrained model
+# Function to generate a comforting story using GPT-2
 def generate_story(theme):
-    story_prompt = f"Tell a detailed, comforting, and heartwarming story about {theme}. Include a character facing a tough challenge, showing courage, and overcoming it with a positive resolution."
-    input_ids = tokenizer.encode(story_prompt, return_tensors='pt')
-    story_ids = model.generate(
+    # 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 GPT-2
+    input_ids = gpt2_tokenizer.encode(story_prompt, return_tensors='pt')
+    
+    story_ids = gpt2_model.generate(
         input_ids,
-        max_length=500,
-        temperature=0.9,
-        repetition_penalty=1.1,
+        max_length=500,  # Generate longer stories
+        temperature=0.8,  # Balanced creativity
+        top_p=0.9,
+        repetition_penalty=1.2,
         num_return_sequences=1
     )
-    story = tokenizer.decode(story_ids[0], skip_special_tokens=True)
+    
+    # Decode the generated text
+    story = gpt2_tokenizer.decode(story_ids[0], skip_special_tokens=True)
     return story
 
+
 # Function to generate an empathetic response
 def generate_response(user_input):
     response_prompt = f"You are a compassionate support bot. A user has shared: '{user_input}'. Respond with empathy and encouragement."
@@ -72,6 +83,8 @@ with st.sidebar:
             tts.save(meditation_audio)
         st.audio(meditation_audio, format="audio/mp3")
 
+# Sidebar for additional features
+with st.sidebar:
     st.header("📖 Short Comforting Story")
     story_theme = st.selectbox("Choose a theme for your story:", ["courage", "healing", "hope"])
     if st.button("Generate Story"):
@@ -79,6 +92,7 @@ with st.sidebar:
             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)