Maximum number of new tokens slider

app.py

@@ -13,12 +13,12 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 n_gpu = 0 if torch.cuda.is_available()==False else torch.cuda.device_count()
 model.to(device)
 
-def generate(text = ""):  
+def generate(text = "", max_new_tokens = 128):  
   streamer = TextIteratorStreamer(tok, timeout=10.)
   if len(text) == 0:
     text = " "
   inputs = tok([text], return_tensors="pt").to(device)
-  generation_kwargs = dict(inputs, streamer=streamer, repetition_penalty=2.0, do_sample=True, top_k=40, top_p=0.97, max_new_tokens=128, pad_token_id = model.config.eos_token_id, early_stopping=True, no_repeat_ngram_size=4)
+  generation_kwargs = dict(inputs, streamer=streamer, repetition_penalty=2.0, do_sample=True, top_k=40, top_p=0.97, max_new_tokens=max_new_tokens, pad_token_id = model.config.eos_token_id, early_stopping=True, no_repeat_ngram_size=4)
   thread = Thread(target=model.generate, kwargs=generation_kwargs)
   thread.start()
   generated_text = ""
@@ -35,8 +35,10 @@ def generate(text = ""):
 demo = gr.Interface(
     title="TextIteratorStreamer + Gradio demo",
     fn=generate,
-    inputs=gr.inputs.Textbox(lines=5, label="Input Text"),
-    outputs=gr.outputs.Textbox(label="Generated Text"),    
+    inputs=[gr.inputs.Textbox(lines=5, label="Input Text"),
+            gr.inputs.Slider(default=128,minimum=5, maximum=256, step=1, label="Maximum number of new tokens")],
+    outputs=gr.outputs.Textbox(label="Generated Text"),
+    allow_flagging="never"
 )
 
 demo.queue()