Update app.py

app.py

@@ -6,8 +6,8 @@ import os
 # --- App Configuration ---
 TITLE = "✍️ AI Story Outliner"
 DESCRIPTION = """
-Enter a prompt and get 10 unique story outlines from a CPU-friendly AI model.
-The app uses **Microsoft's Phi-2**, a powerful small language model, to generate creative outlines.
+Enter a prompt and get 10 unique story outlines from a powerful AI model.
+The app uses **Mistral-7B-v0.1**, a popular and capable open-source model, to generate creative outlines.
 
 **How it works:**
 1.  Enter your story idea.
@@ -25,7 +25,7 @@ examples = [
 ]
 
 # --- Model Initialization ---
-# This section loads the Phi-2 model, which requires authentication.
+# This section loads the Mistral-7B model, which requires authentication.
 # It will automatically use the HF_TOKEN secret when deployed on Hugging Face Spaces.
 generator = None
 model_error = None
@@ -42,16 +42,15 @@ try:
         # If no token is found, raise an error to prevent the app from crashing later.
         raise ValueError("Hugging Face token not found. Please set the HF_TOKEN secret in your Space settings.")
 
-    # Using 'microsoft/phi-2'. This model requires authentication and trust_remote_code=True.
+    # Using 'mistralai/Mistral-7B-v0.1'. This model is gated and requires a token.
     generator = pipeline(
         "text-generation",
-        model="microsoft/phi-2",
+        model="mistralai/Mistral-7B-v0.1",
         token=hf_token,
         torch_dtype=torch.bfloat16, # More performant data type
-        device_map="auto", # Will use GPU if available, otherwise CPU
-        trust_remote_code=True # Required for Phi-2 model
+        device_map="auto" # Will use GPU if available, otherwise CPU
     )
-    print("✅ microsoft/phi-2 model loaded successfully!")
+    print("✅ mistralai/Mistral-7B-v0.1 model loaded successfully!")
 
 except Exception as e:
     model_error = e
@@ -77,10 +76,9 @@ def generate_stories(prompt: str) -> list[str]:
         return [""] * 10
 
     try:
-        # A clear, instructive prompt format that works well with models like Phi-2.
-        story_prompt = f"""Instruct: Create a short story outline based on this idea: "{prompt}"
-The outline should have three parts: a dramatic hook, a concise ballad, and a satisfying finale. Use emojis.
-Output:
+        # This prompt format is optimized for Mistral instruct models.
+        story_prompt = f"""[INST] Create a short story outline based on this idea: "{prompt}"
+The outline must have three parts: a dramatic hook, a concise ballad, and a satisfying finale. Use emojis for each section header. [/INST]
 ### 🎬 The Hook
 """
 
@@ -89,8 +87,7 @@ Output:
             "max_new_tokens": 250,
             "num_return_sequences": 10,
             "do_sample": True,
-            "temperature": 0.9,
-            "top_k": 50,
+            "temperature": 0.8,
             "top_p": 0.95,
             "pad_token_id": generator.tokenizer.eos_token_id
         }
@@ -107,8 +104,9 @@ Output:
             full_text = out['generated_text']
             # Add back the part of the prompt we want to see in the output
             story_start = "### 🎬 The Hook\n"
-            generated_part = full_text.split(story_start)[-1]
-            stories.append(story_start + generated_part)
+            # Split by the instruction closing tag to get only the model's response
+            generated_part = full_text.split("[/INST]")[-1].strip()
+            stories.append(generated_part)
 
         # Ensure we return exactly 10 stories, padding if necessary.
         while len(stories) < 10:
@@ -133,7 +131,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=".gradio-container {max-width: 95% !i
         with gr.Column(scale=1):
             input_area = gr.TextArea(
                 lines=5,
-                label="Your Story Prompt �",
+                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 10 Outlines ✨", variant="primary")
@@ -167,4 +165,4 @@ with gr.Blocks(theme=gr.themes.Soft(), css=".gradio-container {max-width: 95% !i
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()