Update app.py

app.py

@@ -4,128 +4,110 @@ from transformers import pipeline
 import os
 
 # --- App Configuration ---
-title = "📗 Health and Mindful Story Gen ❤️"
-description = """
-Enter a topic or a starting sentence related to health, mindfulness, or well-being.
-The app will generate continuations from three different language models.
-**Note:** These models are very large. Initial loading and first-time generation may take a few minutes.
+TITLE = "✍️ AI Story Weaver"
+DESCRIPTION = """
+Enter a prompt, a topic, or the beginning of a story, and get three different continuations from powerful open-source AI models.
+This app uses:
+- **Mistral-7B-Instruct-v0.2**
+- **Google's Gemma-7B-IT**
+- **Meta's Llama-3-8B-Instruct**
+
+**⚠️ Hardware Warning:** These are very large models. Loading them requires a powerful GPU with significant VRAM (ideally > 24GB).
+The initial loading process may take several minutes. You will also need to install the `accelerate` library: `pip install accelerate`
 """
 
-# Define the examples to show in the interface
+# --- Example Prompts for Storytelling ---
 examples = [
-    ["Mental Body Scan"],
-    ["Stretch, Calm, Breath"],
-    ["Relaxed Seat Breath"],
-    ["Walk Feel"],
-    ["Brain gamification"],
-    ["Alleviating stress"],
-    ["Helping breathing, satisfaction"],
-    ["Relieve Stress, Build Support"],
-    ["The Relaxation Response"],
-    ["Taking Deep Breaths"],
-    ["Delete Not Helpful Thoughts"],
-    ["Strengthen Helpful Thoughts"],
-    ["Reprogram Pain and Stress Reactions"],
-    ["How to Sleep Better and Find Joy"],
-    ["Yoga for deep sleep"],
-    ["Being a Happier and Healthier Person"],
-    ["Relieve chronic pain by"],
-    ["Use Mindfulness to Affect Well Being"],
-    ["Build and Boost Mental Strength"],
-    ["Spending Time Outdoors"],
-    ["Daily Routine Tasks"],
-    ["Eating and Drinking - Find Healthy Nutrition Habits"],
-    ["Drinking - Find Reasons and Cut Back or Quit Entirely"],
-    ["Feel better each day when you awake by"],
-    ["Feel better physically by"],
-    ["Practicing mindfulness each day"],
-    ["Be happier by"],
-    ["Meditation can improve health by"],
-    ["Spending time outdoors helps to"],
-    ["Stress is relieved by quieting your mind, getting exercise and time with nature"],
-    ["Break the cycle of stress and anxiety"],
-    ["Feel calm in stressful situations"],
-    ["Deal with work pressure by"],
-    ["Learn to reduce feelings of being overwhelmed"]
+    ["The old lighthouse keeper stared into the storm. He'd seen many tempests, but this one was different. This one had eyes..."],
+    ["In a city powered by dreams, a young inventor creates a machine that can record them. His first recording reveals a nightmare that doesn't belong to him."],
+    ["The knight adjusted his helmet, the dragon's roar echoing in the valley. He was ready for the fight, but not for what the dragon said when it finally spoke."],
+    ["She found the old leather-bound journal in her grandfather's attic. The first entry read: 'To relieve stress, I walk in the woods. But today, the woods walked with me.'"],
+    ["The meditation app promised to help her 'delete unhelpful thoughts.' She tapped the button, and to her horror, the memory of her own name began to fade..."]
 ]
 
 # --- Model Initialization ---
-# WARNING: Loading these models requires significant hardware (ideally a GPU with >24GB VRAM).
-# 'device_map="auto"' and 'torch_dtype' require the 'accelerate' library.
-# Install dependencies: pip install gradio transformers torch accelerate
+# This section loads the models. It requires significant hardware resources.
+# `device_map="auto"` and `torch_dtype="auto"` help manage resources by using available GPUs and half-precision.
 try:
     print("Initializing models... This may take several minutes.")
     
-    # Using device_map="auto" to automatically use available GPUs.
-    # Using torch_dtype="auto" to load models in half-precision (float16/bfloat16) to save memory.
-    
-    generator1 = pipeline("text-generation", model="gpt2-large", device_map="auto")
-    print("GPT-2 Large loaded.")
+    # NOTE: For Llama-3, you may need to log in to Hugging Face and accept the license agreement.
+    # from huggingface_hub import login
+    # login("YOUR_HF_TOKEN")
+
+    generator1 = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.2", torch_dtype="auto", device_map="auto")
+    print("✅ Mistral-7B loaded.")
     
-    generator2 = pipeline("text-generation", model="EleutherAI/gpt-neo-2.7B", torch_dtype="auto", device_map="auto")
-    print("GPT-Neo 2.7B loaded.")
+    generator2 = pipeline("text-generation", model="google/gemma-7b-it", torch_dtype="auto", device_map="auto")
+    print("✅ Gemma-7B loaded.")
 
-    generator3 = pipeline("text-generation", model="EleutherAI/gpt-j-6B", torch_dtype="auto", device_map="auto")
-    print("GPT-J 6B loaded.")
+    generator3 = pipeline("text-generation", model="meta-llama/Llama-3-8B-Instruct", torch_dtype="auto", device_map="auto")
+    print("✅ Llama-3-8B loaded.")
 
-    print("All models loaded successfully! ✅")
+    print("All models loaded successfully! 🎉")
 
 except Exception as e:
-    print(f"Error loading models: {e}")
-    print("Please ensure you have 'torch' and 'accelerate' installed and have sufficient VRAM.")
-    # Create dummy functions if models fail to load, so the app can still launch.
+    print(f"--- 🚨 Error loading models ---")
+    print(f"Error: {e}")
+    print("Please ensure you have 'torch' and 'accelerate' installed, have sufficient VRAM, and are logged into Hugging Face if required.")
+    # Create a dummy function if models fail, so the app can still launch with an error message.
     def failed_generator(prompt, **kwargs):
-        return [{'generated_text': "Model failed to load. Check console for errors."}]
+        return [{'generated_text': "A model failed to load. Please check the console for errors. You may need more VRAM or need to accept model license terms on Hugging Face."}]
     generator1 = generator2 = generator3 = failed_generator
 
 
 # --- App Logic ---
-def generate_outputs(input_text: str) -> tuple[str, str, str]:
-    """Generates text from the three loaded models."""
-    # Using 'max_new_tokens' is preferred over 'max_length' to specify the length of the generated text only.
-    params = {"max_new_tokens": 60, "num_return_sequences": 1}
+def generate_stories(prompt: str) -> tuple[str, str, str]:
+    """Generates text from the three loaded models based on the user's prompt."""
+    if not prompt:
+        return "Please enter a prompt to start.", "", ""
+        
+    # We use 'max_new_tokens' to control the length of the generated story.
+    # Increased to 200 for more substantial story continuations.
+    params = {"max_new_tokens": 200, "do_sample": True, "temperature": 0.7, "top_p": 0.95}
     
-    out1 = generator1(input_text, **params)[0]['generated_text']
-    out2 = generator2(input_text, **params)[0]['generated_text']
-    out3 = generator3(input_text, **params)[0]['generated_text']
+    # Generate from all three models
+    out1 = generator1(prompt, **params)[0]['generated_text']
+    out2 = generator2(prompt, **params)[0]['generated_text']
+    out3 = generator3(prompt, **params)[0]['generated_text']
     
     return out1, out2, out3
 
 # --- Gradio Interface ---
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown(f"<h1 style='text-align: center;'>{title}</h1>")
-    gr.Markdown(description)
+with gr.Blocks(theme=gr.themes.Soft(), css=".gradio-container {max-width: 95% !important;}") as demo:
+    gr.Markdown(f"<h1 style='text-align: center;'>{TITLE}</h1>")
+    gr.Markdown(DESCRIPTION)
 
     with gr.Row():
         with gr.Column(scale=1):
             input_area = gr.TextArea(
-                lines=3,
-                label="Your starting prompt 👇",
-                placeholder="e.g., 'To relieve stress, I will try...'"
+                lines=5,
+                label="Your Story Prompt 👇",
+                placeholder="e.g., 'The last dragon on Earth lived not in a cave, but in a library...'"
             )
-            generate_button = gr.Button("Generate ✨", variant="primary")
+            generate_button = gr.Button("Weave a Story ✨", variant="primary")
     
         with gr.Column(scale=2):
             with gr.Tabs():
-                with gr.TabItem("GPT-2 Large"):
-                    gen1_output = gr.TextArea(label="GPT-2 Large Output", interactive=False, lines=7)
-                with gr.TabItem("GPT-Neo 2.7B"):
-                    gen2_output = gr.TextArea(label="GPT-Neo 2.7B Output", interactive=False, lines=7)
-                with gr.TabItem("GPT-J 6B"):
-                    gen3_output = gr.TextArea(label="GPT-J 6B Output", interactive=False, lines=7)
+                with gr.TabItem("Mistral-7B"):
+                    gen1_output = gr.TextArea(label="Mistral's Tale", interactive=False, lines=12)
+                with gr.TabItem("Gemma-7B"):
+                    gen2_output = gr.TextArea(label="Gemma's Chronicle", interactive=False, lines=12)
+                with gr.TabItem("Llama-3-8B"):
+                    gen3_output = gr.TextArea(label="Llama's Legend", interactive=False, lines=12)
 
     gr.Examples(
         examples=examples,
         inputs=input_area,
-        label="Example Prompts (Click to use)"
+        label="Example Story Starters (Click to use)"
     )
 
     generate_button.click(
-        fn=generate_outputs,
+        fn=generate_stories,
         inputs=input_area,
         outputs=[gen1_output, gen2_output, gen3_output],
         api_name="generate"
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()