Update app.py

app.py

@@ -24,6 +24,9 @@ inputs = [
     gr.Textbox(label="Main Character", placeholder="e.g. a brave knight"),
     gr.Textbox(label="Setting", placeholder="e.g. in an enchanted forest"),
     gr.Textbox(label="Plot Twist", placeholder="e.g. discovers a hidden treasure")
+    gr.Slider(minimum=1, maximum=2048, value=64, step=1, label="Max new tokens"),
+    gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
+    gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)")
 ]
 
 outputs = gr.Textbox(label="Generated Story")
@@ -54,12 +57,12 @@ def respond(message, history, system_message, max_tokens, temperature, top_p):
     return output_text 
 
 @spaces.GPU 
-def generate_story(character, setting, plot_twist):
+def generate_story(character, setting, plot_twist, max_tokens, temperature, top_p):
     """Define the function to generate the story."""
     prompt = f"Write a short story with the following details:\nMain character: {character}\nSetting: {setting}\nPlot twist: {plot_twist}\n\nStory:"
     input_ids = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
 
-    output_ids = model.generate(input_ids, max_length=50, num_return_sequences=1)
+    output_ids = model.generate(input_ids, max_length=max_tokens, num_return_sequences=1, temperature=temperature, top_p=top_p)
 
     output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)