Update app.py

app.py

@@ -589,54 +589,79 @@ demo.launch()'''
 
   
 import gradio as gr
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-# Load a lightweight model (distilgpt2) for initial story generation
-generator = pipeline("text-generation", model="distilgpt2", max_length=50, truncation=True)
+# Load TinyLlama model and tokenizer
+model_name = "TinyLlama/TinyLlama-1.1B-Chat"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
 
 # Initialize story state
 story_history = []
 
+def generate_text(prompt, max_length=100):
+    """Generate text using TinyLlama model with optimized settings."""
+    # Format prompt for chat
+    formatted_prompt = f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
+    
+    # Tokenize input
+    inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
+    
+    # Generate with optimized parameters
+    with torch.no_grad():
+        outputs = model.generate(
+            inputs.input_ids,
+            max_new_tokens=max_length,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9,
+            num_return_sequences=1,
+            pad_token_id=tokenizer.eos_token_id
+        )
+    
+    # Decode and clean up response
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    response = response.replace(formatted_prompt, "").strip()
+    return response
+
 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."
+        f"Write a short, engaging paragraph (40-50 words) for a story starting with: '{prompt}'. "
+        "Make it descriptive and ensure it relates directly to the prompt."
     )
-    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_output = generate_text(story_prompt, max_length=75)
     
-    # Ensure prompt relevance
-    prompt_words = set(prompt.lower().split())
-    if not any(word in story_output.lower() for word in prompt_words):
+    # Ensure the story starts with the prompt theme
+    if not any(word in story_output.lower() for word in prompt.lower().split()):
         story_output = f"{prompt}. {story_output}"
     
     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")
+def generate_options(story_so_far):
+    """Generate three context-aware options based on the current story."""
+    option_prompt = (
+        f"Based on this story: '{story_so_far}'\n"
+        "Generate 3 distinct, exciting options for what happens next. "
+        "Each option should be a single sentence and directly relate to the story's context. "
+        "Format: 1. [Option 1] 2. [Option 2] 3. [Option 3]"
+    )
     
-    # 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_text = generate_text(option_prompt, max_length=100)
+    
+    # Parse options or provide fallback options if generation fails
+    try:
+        options = [opt.split(". ")[1].strip() for opt in options_text.split("\n") if opt.strip().startswith(("1.", "2.", "3."))]
+        if len(options) != 3:
+            raise ValueError("Invalid options generated")
+    except:
+        # Fallback options based on common story elements
+        character = story_so_far.split()[0]
         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."
+            f"{character} faces an unexpected challenge that tests their resolve.",
+            f"{character} discovers something that changes everything.",
+            f"{character} makes a decision that leads to adventure."
         ]
     
     return options
@@ -645,26 +670,39 @@ def start_story(initial_prompt):
     """Start the story with the user's initial prompt."""
     global story_history
     story_history = [initial_prompt]
+    
+    # Generate initial segment
     segment = generate_initial_story(initial_prompt)
     story_history.append(segment)
+    
+    # Generate options
     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]
+    
+    return format_story_output(segment, options), *options
 
 def continue_story(choice):
     """Continue the story based on the user's choice."""
     global story_history
-    story_history.append(f"You chose: {choice}")
-    # 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_history.append(f"Choice: {choice}")
+    
+    # Generate new segment
+    continuation_prompt = (
+        f"Continue this story coherently:\n{story_history[-2]}\n"
+        f"The chosen path: {choice}\n"
+        "Write a natural 40-50 word continuation that builds on these events."
     )
-    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]
+    new_segment = generate_text(continuation_prompt, max_length=75)
+    story_history.append(new_segment)
+    
+    # Generate new options
+    options = generate_options(new_segment)
+    
+    return format_story_output("\n".join(story_history), options), *options
+
+def format_story_output(story, options):
+    """Format the story and options for display."""
+    return f"{story}\n\n🔹 What happens next?\n1️⃣ {options[0]}\n2️⃣ {options[1]}\n3️⃣ {options[2]}"
 
 def reset_story():
     """Reset the story to start fresh."""
@@ -674,26 +712,38 @@ def reset_story():
 
 # Gradio interface
 with gr.Blocks(title="Story Adventure") as demo:
-    gr.Markdown("# Story Adventure")
-    gr.Markdown("Begin your tale and shape it with your choices!")
+    gr.Markdown("# 📚 Story Adventure")
+    gr.Markdown("Create your own adventure! Enter a prompt and shape the story through your choices.")
     
     with gr.Row():
-        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")
+        prompt_input = gr.Textbox(
+            label="Story Prompt",
+            placeholder="e.g., 'Sarah discovers a mysterious letter in her attic'",
+            lines=2
+        )
+        start_button = gr.Button("🎮 Start Story", variant="primary")
+    
+    story_output = gr.Textbox(label="Your Story", lines=12, interactive=False)
     
-    story_output = gr.Textbox(label="Your Story", lines=10, interactive=False)
-    choice_button_1 = gr.Button(value="", visible=False)
-    choice_button_2 = gr.Button(value="", visible=False)
-    choice_button_3 = gr.Button(value="", visible=False)
+    with gr.Row():
+        choice_button_1 = gr.Button(value="", visible=False, variant="secondary")
+        choice_button_2 = gr.Button(value="", visible=False, variant="secondary")
+        choice_button_3 = gr.Button(value="", visible=False, variant="secondary")
     
-    reset_button = gr.Button("Reset Story")
+    reset_button = gr.Button("🔄 Reset Story", variant="stop")
     
+    # Event handlers
     start_button.click(
         fn=start_story,
         inputs=prompt_input,
         outputs=[story_output, choice_button_1, choice_button_2, choice_button_3]
     ).then(
-        fn=lambda story, c1, c2, c3: (story, gr.update(value=c1, visible=True), gr.update(value=c2, visible=True), gr.update(value=c3, visible=True)),
+        fn=lambda story, c1, c2, c3: (
+            story,
+            gr.update(value=c1, visible=True),
+            gr.update(value=c2, visible=True),
+            gr.update(value=c3, visible=True)
+        ),
         inputs=[story_output, choice_button_1, choice_button_2, choice_button_3],
         outputs=[story_output, choice_button_1, choice_button_2, choice_button_3]
     )
@@ -704,7 +754,12 @@ with gr.Blocks(title="Story Adventure") as demo:
             inputs=button,
             outputs=[story_output, choice_button_1, choice_button_2, choice_button_3]
         ).then(
-            fn=lambda story, c1, c2, c3: (story, gr.update(value=c1, visible=True), gr.update(value=c2, visible=True), gr.update(value=c3, visible=True)),
+            fn=lambda story, c1, c2, c3: (
+                story,
+                gr.update(value=c1, visible=True),
+                gr.update(value=c2, visible=True),
+                gr.update(value=c3, visible=True)
+            ),
             inputs=[story_output, choice_button_1, choice_button_2, choice_button_3],
             outputs=[story_output, choice_button_1, choice_button_2, choice_button_3]
         )
@@ -714,4 +769,5 @@ with gr.Blocks(title="Story Adventure") as demo:
         outputs=[story_output, choice_button_1, choice_button_2, choice_button_3]
     )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()