Update app.py

app.py

@@ -15,7 +15,7 @@ def load_model(model_name):
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(model_name)
     return model, tokenizer
-def extend(input_text, max_size=20, top_k=50, top_p=0.95):
+def extend(input_text, max_size=20, top_k=50, top_p=0.95, num_return_sequences):
     if len(input_text) == 0:
         input_text = ""
     encoded_prompt = tokenizer.encode(
@@ -32,7 +32,7 @@ def extend(input_text, max_size=20, top_k=50, top_p=0.95):
     top_k=top_k, 
     top_p=top_p, 
     do_sample=True,
-    num_return_sequences=1)
+    num_return_sequences=num_return_sequences)
     # Remove the batch dimension when returning multiple sequences
     if len(output_sequences.shape) > 2:
         output_sequences.squeeze_()