Update app.py

app.py

@@ -1,65 +1,116 @@
 import gradio as gr
 from transformers import pipeline
 
-# Load the TinyLlama model for text generation
+# Load TinyLlama model for text generation
 generator = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
 
-# Store the current story in a global variable
+# Store the current story and last options in global variables
 current_story = ""
+last_options = []
 
 def start_story(prompt):
-    global current_story
-    # Generate the initial story segment, ensuring the prompt is the starting point
-    result = generator(
-        prompt,
-        max_new_tokens=150,  # Increased to ~100-150 words (approx. 10 lines)
-        temperature=0.7,     # Lower temperature for less randomness
-        top_k=50,            # Filter to top 50 likely next words
-        do_sample=True,      # Enable sampling for creativity
+    global current_story, last_options
+    # Generate the initial story segment
+    formatted_prompt = f"Human: Write a creative story starting with: '{prompt}' in about 100-150 words."
+    story_result = generator(
+        formatted_prompt,
+        max_new_tokens=150,  # ~10 lines (100-150 words)
+        temperature=0.7,
+        top_k=50,
+        do_sample=True,
         truncation=True
     )[0]["generated_text"]
-    current_story = result
-    return current_story
+    current_story = story_result.replace(formatted_prompt, "").strip()
+    
+    # Generate story-specific options
+    options_prompt = f"Human: Based on this story: '{current_story}', suggest three distinct options for what happens next. Each option should be 10-20 words and continue the narrative logically."
+    options_result = generator(
+        options_prompt,
+        max_new_tokens=100,  # Enough for 3 options
+        temperature=0.8,     # Slightly higher for variety
+        top_k=50,
+        do_sample=True,
+        truncation=True
+    )[0]["generated_text"]
+    options_text = options_result.replace(options_prompt, "").strip()
+    
+    # Parse options, with fallback if generation is messy
+    options_lines = [line.strip() for line in options_text.split("\n") if line.strip()]
+    last_options = options_lines[:3] if len(options_lines) >= 3 else [
+        "The story took an unexpected twist.",
+        "A new challenge emerged from the shadows.",
+        "The journey led to an unlikely ally."
+    ]
+    
+    options_display = "\n\n**Choose an option:**\n" + "\n".join([f"{i+1}. {opt}" for i, opt in enumerate(last_options)])
+    return current_story + options_display
 
-def continue_story():
-    global current_story
+def continue_story(option_num):
+    global current_story, last_options
     if not current_story:
         return "Please start a story first!"
-    # Continue the story, using the current story as the prompt
-    result = generator(
-        current_story,
-        max_new_tokens=150,  # Increased to ~100-150 words (approx. 10 lines)
+    if not last_options or not (1 <= int(option_num) <= 3):
+        return "Invalid option! Choose 1, 2, or 3."
+    
+    # Continue the story with the chosen option
+    chosen_option = last_options[int(option_num) - 1]
+    formatted_prompt = f"Human: Continue this story: '{current_story}' with this direction: '{chosen_option}' in about 100-150 words."
+    story_result = generator(
+        formatted_prompt,
+        max_new_tokens=150,  # ~10 lines
         temperature=0.7,
         top_k=50,
         do_sample=True,
         truncation=True
     )[0]["generated_text"]
-    current_story = result
-    return current_story
+    current_story = story_result.replace(formatted_prompt, "").strip()
+    
+    # Generate new story-specific options
+    options_prompt = f"Human: Based on this story: '{current_story}', suggest three distinct options for what happens next. Each option should be 10-20 words and continue the narrative logically."
+    options_result = generator(
+        options_prompt,
+        max_new_tokens=100,
+        temperature=0.8,
+        top_k=50,
+        do_sample=True,
+        truncation=True
+    )[0]["generated_text"]
+    options_text = options_result.replace(options_prompt, "").strip()
+    
+    options_lines = [line.strip() for line in options_text.split("\n") if line.strip()]
+    last_options = options_lines[:3] if len(options_lines) >= 3 else [
+        "The story took an unexpected twist.",
+        "A new challenge emerged from the shadows.",
+        "The journey led to an unlikely ally."
+    ]
+    
+    options_display = "\n\n**Choose an option:**\n" + "\n".join([f"{i+1}. {opt}" for i, opt in enumerate(last_options)])
+    return current_story + options_display
 
 def reset_story():
-    global current_story
-    # Reset the story to empty
+    global current_story, last_options
     current_story = ""
+    last_options = []
     return "Story reset. Enter a new prompt to begin!"
 
 # Create the Gradio interface
-with gr.Blocks(title="AI Story Generator") as demo:
-    gr.Markdown("# AI Story Generator")
-    gr.Markdown("Enter a prompt to start your story, then continue or reset it!")
+with gr.Blocks(title="AI Story Game") as demo:
+    gr.Markdown("# AI Story Game")
+    gr.Markdown("Start your adventure with a prompt, then guide the story with options or reset it!")
     
     # Input for the initial prompt
-    prompt_input = gr.Textbox(label="Start your story with a prompt", placeholder="E.g., 'A knight rode into a dark forest...'")
-    output = gr.Textbox(label="Your Story", lines=15)  # Increased lines to fit longer output
+    prompt_input = gr.Textbox(label="Start your story with a prompt", placeholder="E.g., 'A thief crept through the shadows...'")
+    output = gr.Textbox(label="Your Story", lines=15)
     
-    # Buttons for starting, continuing, and resetting
+    # Option input and buttons
+    option_input = gr.Textbox(label="Enter option number (1-3)", placeholder="E.g., 1", lines=1)
     start_button = gr.Button("Start Story")
     continue_button = gr.Button("Continue Story")
     reset_button = gr.Button("Reset Story")
     
     # Connect buttons to functions
     start_button.click(fn=start_story, inputs=prompt_input, outputs=output)
-    continue_button.click(fn=continue_story, inputs=None, outputs=output)
+    continue_button.click(fn=continue_story, inputs=option_input, outputs=output)
     reset_button.click(fn=reset_story, inputs=None, outputs=output)
 
 # Launch the app