Update app.py

app.py

@@ -1,5 +1,5 @@
 
-'''import gradio as gr
+import gradio as gr
 from transformers import pipeline
 
 # Load the gpt-neo-125M model for text generation
@@ -13,7 +13,7 @@ def start_story(prompt):
     # Generate the initial story segment, ensuring the prompt is the starting point
     result = generator(
         prompt,
-        max_new_tokens=50,  # Limit new tokens to keep it concise
+        max_new_tokens=150,  # Limit new tokens to keep it concise
         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
@@ -29,7 +29,7 @@ def continue_story():
     # Continue the story, using the current story as the prompt
     result = generator(
         current_story,
-        max_new_tokens=50,
+        max_new_tokens=150,
         temperature=0.7,
         top_k=50,
         do_sample=True,
@@ -63,128 +63,5 @@ with gr.Blocks(title="AI Story Generator") as demo:
     continue_button.click(fn=continue_story, inputs=None, outputs=output)
     reset_button.click(fn=reset_story, inputs=None, outputs=output)
 
-# Launch the app
-demo.launch()'''
-
-
-import gradio as gr
-from transformers import pipeline
-
-# Load TinyLlama model for text generation
-generator = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
-
-# Store the current story and last options in global variables
-current_story = ""
-last_options = []
-
-def start_story(prompt):
-    global current_story, last_options
-    # Format the prompt for TinyLlama to generate a detailed story
-    formatted_prompt = f"Human: Write a creative story starting with: '{prompt}' in about 100-150 words."
-    # Generate a longer initial story segment
-    result = generator(
-        formatted_prompt,
-        max_new_tokens=150,  # ~100-150 words
-        temperature=0.7,     # Controlled creativity
-        top_k=50,            # Coherence
-        do_sample=True,
-        truncation=True
-    )[0]["generated_text"]
-    # Strip the "Human:" prefix from the output
-    story_text = result.replace(formatted_prompt, "").strip()
-    current_story = story_text
-    
-    # Generate dynamic options based on the prompt
-    last_options = generate_options(prompt)
-    options_text = "\n\n**Choose an option:**\n" + "\n".join([f"{i+1}. {opt}" for i, opt in enumerate(last_options)])
-    return current_story + options_text
-
-def continue_story(option_num):
-    global current_story, last_options
-    if not current_story:
-        return "Please start a story first!"
-    if not last_options or not (1 <= int(option_num) <= 3):
-        return "Invalid option or no story started yet!"
-    
-    # Use the selected option to guide the next segment
-    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."
-    result = generator(
-        formatted_prompt,
-        max_new_tokens=150,
-        temperature=0.7,
-        top_k=50,
-        do_sample=True,
-        truncation=True
-    )[0]["generated_text"]
-    # Update the story
-    story_text = result.replace(formatted_prompt, "").strip()
-    current_story = story_text
-    
-    # Generate new options based on the updated story
-    last_options = generate_options(current_story)
-    options_text = "\n\n**Choose an option:**\n" + "\n".join([f"{i+1}. {opt}" for i, opt in enumerate(last_options)])
-    return current_story + options_text
-
-def reset_story():
-    global current_story, last_options
-    current_story = ""
-    last_options = []
-    return "Story reset. Enter a new prompt to begin!"
-
-def generate_options(text):
-    # Dynamic option generation based on keywords in the prompt or story
-    text_lower = text.lower()
-    if "girl" in text_lower or "woman" in text_lower or "she" in text_lower:
-        return [
-            "She decided to investigate the mystery further.",
-            "She encountered someone unexpected who changed her path.",
-            "She chose to keep her discovery a secret for now."
-        ]
-    elif "man" in text_lower or "he" in text_lower or "boy" in text_lower:
-        return [
-            "He set out to confront the source of the mystery.",
-            "He stumbled upon a hidden truth that tested his resolve.",
-            "He decided to seek help from an old ally."
-        ]
-    elif "book" in text_lower or "library" in text_lower or "knowledge" in text_lower:
-        return [
-            "The knowledge within led to a surprising revelation.",
-            "A stranger appeared, claiming ownership of the discovery.",
-            "The pursuit of answers took a dangerous turn."
-        ]
-    elif "cat" in text_lower or "dog" in text_lower or "animal" in text_lower:
-        return [
-            "It followed a strange sound into the night.",
-            "It met a companion that altered its journey.",
-            "It uncovered a hidden clue beneath the surface."
-        ]
-    else:
-        return [
-            "The situation escalated into an unexpected conflict.",
-            "A new clue emerged, pointing to a distant location.",
-            "The choice was made to abandon the path and start anew."
-        ]
-
-# Create the Gradio interface
-with gr.Blocks(title="AI Story Generator") as demo:
-    gr.Markdown("# AI Story Generator")
-    gr.Markdown("Enter any prompt to start your story, then choose an option to continue or reset!")
-    
-    # Input for the initial prompt
-    prompt_input = gr.Textbox(label="Start your story with a prompt", placeholder="E.g., 'A cat sat on a rooftop...'")
-    output = gr.Textbox(label="Your Story", lines=15)
-    
-    # Buttons and option input
-    start_button = gr.Button("Start Story")
-    option_input = gr.Textbox(label="Enter option number (1-3)", placeholder="E.g., 1", lines=1)
-    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=option_input, outputs=output)
-    reset_button.click(fn=reset_story, inputs=None, outputs=output)
-
 # Launch the app
 demo.launch()