Update app.py

app.py

@@ -5,9 +5,9 @@ from transformers import pipeline
 text_generator = pipeline("text-generation", model="gpt2")
 
 # Define the function for story generation
-def generate_story(prompt):
-    # Generate a story based on the user's prompt
-    generated_text = text_generator(prompt, max_length=300, num_return_sequences=1)[0]['generated_text']
+def generate_story(prompt, word_count):
+    # Generate a story based on the user's prompt and word count
+    generated_text = text_generator(prompt, max_length=word_count, num_return_sequences=1)[0]['generated_text']
     return generated_text
 
 # Define example inputs for the Gradio interface
@@ -20,13 +20,16 @@ example_inputs = [
     "Deep in the ocean, an underwater explorer encountered a mysterious and ancient creature."
 ]
 
-# Create a Gradio interface with examples
+# Create a Gradio interface with examples and a word count slider
 iface = gr.Interface(
     fn=generate_story,
-    inputs="text",
+    inputs=[
+        gr.components.Textbox(label="Prompt"),
+        gr.components.Slider(minimum=50, maximum=500, default=100, label="Word Count")
+    ],
     outputs="text",
-    title="Story Generator",
-    description="Enter a prompt to generate a story.",
+    title="Story Generator with Word Count",
+    description="Enter a prompt and select the word count to generate a story.",
     examples=example_inputs
 )