Update app.py

app.py

@@ -13,7 +13,7 @@ def correct_text(text, genConfig):
     corrected_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return corrected_text
 
-def respond(text, max_new_tokens, min_new_tokens, num_beams, num_beam_groups, temperature, top_k, top_p, guidance_scale):
+def respond(text, max_new_tokens, min_new_tokens, num_beams, num_beam_groups, temperature, top_k, top_p, no_repeat_ngram_size, guidance_scale):
     config = GenerationConfig(
         max_new_tokens=max_new_tokens,
         min_new_tokens=min_new_tokens,
@@ -22,6 +22,7 @@ def respond(text, max_new_tokens, min_new_tokens, num_beams, num_beam_groups, te
         temperature=float(temperature),
         top_k=top_k,
         top_p=float(top_p),
+        no_repeat_ngram_size=no_repeat_ngram_size,
         early_stopping=True,
         do_sample=True
     )
@@ -61,7 +62,9 @@ with gr.Blocks() as demo:
         top_k       = gr.Slider(minimum=0,   maximum=200,   value=50,  step=1,    label="Top-k")
         top_p       = gr.Slider(minimum=0.1, maximum=1.0,   value=1.0, step=0.05, label="Top-p (nucleus sampling)")
         guideScale  = gr.Slider(minimum=0.1, maximum=50.0,  value=1.0, step=1, label="Guidance Scale")
+        no_repeat_ngram_size = gr.Slider(0, 20, 0, 1, "Limit N-grams of given Size")
+        do_sample = gr.Checkbox(value=True, label="Do Sampling")
         
-    submitBtn.click(respond, [prompt_box, max_tokens, min_tokens, num_beams, beam_groups, temperature, top_k, top_p, guideScale], output_box)
+    submitBtn.click(respond, [prompt_box, max_tokens, min_tokens, num_beams, beam_groups, temperature, top_k, top_p, no_repeat_ngram_size, guideScale, do_sample], output_box)
 
 demo.launch()