Update app.py

app.py

@@ -586,70 +586,85 @@ demo.launch()'''
 
 
  
+
+  
 import gradio as gr
 from transformers import pipeline
-import random
 
-# Load a lighter model (distilgpt2 is faster than TinyLlama)
+# Load a lightweight model (distilgpt2) for initial story generation
 generator = pipeline("text-generation", model="distilgpt2", max_length=50, truncation=True)
 
 # Initialize story state
 story_history = []
 
-def generate_story_segment(prompt, user_choice=None):
-    """Generate a short story segment and three specific options in one go."""
-    if user_choice:
-        prompt = f"{prompt} {user_choice}"
-    
-    # Single prompt for both story and choices
-    combined_prompt = (
-        f"Write a short paragraph (20-30 words) based on this: '{prompt}'. "
-        "Then, list three short, specific options (1 sentence each) to expand the story logically."
+def generate_initial_story(prompt):
+    """Generate a coherent initial story paragraph from the prompt."""
+    story_prompt = (
+        f"Write a short paragraph (30-40 words) introducing a story based on: '{prompt}'. "
+        "Keep it simple, coherent, and descriptive."
     )
-    output = generator(combined_prompt, max_length=100, num_return_sequences=1, do_sample=True, temperature=0.7)[0]["generated_text"]
-    output = output.replace(combined_prompt, "").strip()
-
-    # Split story and choices (assuming model outputs them sequentially)
-    lines = output.split("\n")
-    story_output = lines[0] if lines else prompt  # Default to prompt if no story
-    choices = lines[1:4] if len(lines) > 1 else []
-
-    # Ensure the prompt is reflected
+    story_output = generator(story_prompt, max_length=50, num_return_sequences=1, do_sample=True, temperature=0.7)[0]["generated_text"]
+    story_output = story_output.replace(story_prompt, "").strip()
+    
+    # Ensure prompt relevance
     prompt_words = set(prompt.lower().split())
     if not any(word in story_output.lower() for word in prompt_words):
         story_output = f"{prompt}. {story_output}"
-
-    # Refine choices
-    cleaned_choices = []
-    for choice in choices:
-        if len(choice) > 0 and len(choice.split()) > 3:  # Ensure it’s a sentence
-            cleaned_choices.append(choice[:50] + "..." if len(choice) > 50 else choice)
     
-    # Fallback options if not enough choices
-    while len(cleaned_choices) < 3:
-        cleaned_choices.append(random.choice([
-            f"{prompt.split()[0]} meets someone new.",
-            f"{prompt.split()[0]} discovers a hidden secret.",
-            f"{prompt.split()[0]} decides to leave home."
-        ]))
+    return story_output
+
+def generate_options(story):
+    """Generate three specific, relevant options based on the story."""
+    # Extract key elements from the story for relevance
+    story_lower = story.lower()
+    character = story.split()[0]  # Assume first word is the character (e.g., "Priya")
+    
+    # Define option templates based on common story beats
+    if "town" in story_lower or "village" in story_lower:
+        options = [
+            f"{character} moves to a big city for a new job opportunity.",
+            f"{character} meets someone special who changes her life.",
+            f"{character} discovers a family secret that shocks her."
+        ]
+    elif "city" in story_lower:
+        options = [
+            f"{character} lands a dream job that tests her skills.",
+            f"{character} encounters a mysterious stranger in the crowd.",
+            f"{character} finds an old letter in her apartment."
+        ]
+    else:
+        options = [
+            f"{character} embarks on an unexpected adventure.",
+            f"{character} meets a friend who shares a bold idea.",
+            f"{character} uncovers a hidden truth about her past."
+        ]
     
-    return story_output, cleaned_choices[:3]
+    return options
 
 def start_story(initial_prompt):
     """Start the story with the user's initial prompt."""
     global story_history
     story_history = [initial_prompt]
-    segment, choices = generate_story_segment(initial_prompt)
+    segment = generate_initial_story(initial_prompt)
     story_history.append(segment)
-    return "\n".join(story_history), choices[0], choices[1], choices[2]
+    options = generate_options(segment)
+    return f"{segment}\n\n🔹 What happens next? Choose an option:\n1️⃣ {options[0]}\n2️⃣ {options[1]}\n3️⃣ {options[2]}", options[0], options[1], options[2]
 
 def continue_story(choice):
     """Continue the story based on the user's choice."""
     global story_history
-    segment, choices = generate_story_segment(story_history[-1], choice)
     story_history.append(f"You chose: {choice}")
-    story_history.append(segment)
-    return "\n".join(story_history), choices[0], choices[1], choices[2]
+    # Generate a new segment based on the choice
+    story_prompt = (
+        f"Continue this story: '{story_history[-1]} {choice}'. "
+        "Write a short paragraph (30-40 words) that builds on it logically."
+    )
+    story_output = generator(story_prompt, max_length=50, num_return_sequences=1, do_sample=True, temperature=0.7)[0]["generated_text"]
+    story_output = story_output.replace(story_prompt, "").strip()
+    
+    story_history.append(story_output)
+    options = generate_options(story_output)
+    return f"{'\\n'.join(story_history)}\n\n🔹 What happens next? Choose an option:\n1️⃣ {options[0]}\n2️⃣ {options[1]}\n3️⃣ {options[2]}", options[0], options[1], options[2]
 
 def reset_story():
     """Reset the story to start fresh."""
@@ -658,12 +673,12 @@ def reset_story():
     return "Story reset. Enter a new prompt to begin!", "", "", ""
 
 # Gradio interface
-with gr.Blocks(title="Story Builder") as demo:
-    gr.Markdown("# Story Builder")
-    gr.Markdown("Start with any prompt, and build your story step by step!")
+with gr.Blocks(title="Story Adventure") as demo:
+    gr.Markdown("# Story Adventure")
+    gr.Markdown("Begin your tale and shape it with your choices!")
     
     with gr.Row():
-        prompt_input = gr.Textbox(label="Enter your story prompt (e.g., 'There is a boy named Sam')", placeholder="Type here...")
+        prompt_input = gr.Textbox(label="Enter your story prompt (e.g., 'Priya lives in a small town')", placeholder="Type here...")
         start_button = gr.Button("Start Story")
     
     story_output = gr.Textbox(label="Your Story", lines=10, interactive=False)