Update app.py

app.py

@@ -7,7 +7,7 @@ import os
 TITLE = "✍️ AI Story Outliner"
 DESCRIPTION = """
 Enter a prompt and get 10 unique story outlines from a CPU-friendly AI model.
-The app uses **GPT-2**, a reliable and well-established model, to generate creative outlines.
+The app uses **Microsoft's Phi-2**, a powerful small language model, to generate creative outlines.
 
 **How it works:**
 1.  Enter your story idea.
@@ -25,21 +25,33 @@ examples = [
 ]
 
 # --- Model Initialization ---
-# This section loads a smaller, stable, and CPU-friendly model that requires no authentication.
+# This section loads the Phi-2 model, which requires authentication.
+# It will automatically use the HF_TOKEN secret when deployed on Hugging Face Spaces.
 generator = None
 model_error = None
 
 try:
     print("Initializing model... This may take a moment.")
     
-    # Using 'gpt2', a stable and widely supported model that does not require a token.
+    # Explicitly load the token from environment variables (for HF Spaces secrets).
+    hf_token = os.environ.get("HF_TOKEN")
+
+    if hf_token:
+        print("✅ HF_TOKEN secret found.")
+    else:
+        # 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.
     generator = pipeline(
         "text-generation",
-        model="gpt2",
-        torch_dtype=torch.float32, # Use float32 for wider CPU compatibility
-        device_map="auto" # Will use GPU if available, otherwise CPU
+        model="microsoft/phi-2",
+        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
     )
-    print("✅ GPT-2 model loaded successfully!")
+    print("✅ microsoft/phi-2 model loaded successfully!")
 
 except Exception as e:
     model_error = e
@@ -61,21 +73,11 @@ def generate_stories(prompt: str) -> list[str]:
         # Return a list of 10 empty strings to clear the outputs
         return [""] * 10
 
-    # A generic story prompt that works well with models like GPT-2.
-    story_prompt = f"""
-Story Idea: "{prompt}"
-
-Create a short story outline based on this idea.
-
+    # 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:
 ### 🎬 The Hook
-A dramatic opening.
-
-### 🎼 The Ballad
-The main story, told concisely.
-
-### 🔚 The Finale
-A clear and satisfying ending.
----
 """
 
     # Parameters for the pipeline to generate 10 diverse results.
@@ -92,11 +94,15 @@ A clear and satisfying ending.
     # Generate 10 different story variations
     outputs = generator(story_prompt, **params)
 
-    # Extract the generated text. GPT-2 will continue from the prompt.
+    # Extract the generated text.
     stories = []
     for out in outputs:
+        # The model will generate the prompt plus the continuation. We extract just the new part.
         full_text = out['generated_text']
-        stories.append(full_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)
 
     # Ensure we return exactly 10 stories, padding if necessary.
     while len(stories) < 10:
@@ -113,7 +119,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")
@@ -147,4 +153,4 @@ with gr.Blocks(theme=gr.themes.Soft(), css=".gradio-container {max-width: 95% !i
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()