chore: use default generate function

app.py

@@ -31,19 +31,31 @@ def generate(
     temperature=1.0,
     top_p=0.95,
     top_k=20,
+    no_repeat_ngram_size=3,
+    num_beams=1,
 ):
     if input_text.strip() == "":
         return ""
 
     inputs = tokenizer(input_text, return_tensors="pt", add_special_tokens=False)
-    generated = model.custom_generate(
+    # generated = model.custom_generate(
+    #     **inputs,
+    #     parallel_compute_prompt=True,
+    #     max_new_tokens=max_new_tokens,
+    #     do_sample=do_sample,
+    #     temperature=temperature,
+    #     top_p=top_p,
+    #     top_k=top_k,
+    # )
+    generated = model.generate(
         **inputs,
-        parallel_compute_prompt=True,
         max_new_tokens=max_new_tokens,
         do_sample=do_sample,
         temperature=temperature,
         top_p=top_p,
         top_k=top_k,
+        no_repeat_ngram_size=no_repeat_ngram_size,
+        num_beams=num_beams,
     )
     return tokenizer.batch_decode(generated)[0]
 
@@ -97,7 +109,7 @@ with gr.Blocks() as demo:
             label="Max tokens",
             minimum=8,
             maximum=512,
-            value=128,
+            value=64,
             step=4,
         )
         do_sample = gr.Checkbox(
@@ -125,6 +137,20 @@ with gr.Blocks() as demo:
             value=20,
             step=1,
         )
+        no_repeat_ngram_size = gr.Slider(
+            label="No repeat ngram size",
+            minimum=0,
+            maximum=10,
+            value=3,
+            step=1,
+        )
+        num_beams = gr.Slider(
+            label="Num beams",
+            minimum=1,
+            maximum=8,
+            value=1,
+            step=1,
+        )
 
     gr.Examples(
         examples=EXAMPLE_INPUTS,
@@ -140,6 +166,8 @@ with gr.Blocks() as demo:
             temperature,
             top_p,
             top_k,
+            no_repeat_ngram_size,
+            num_beams,
         ],
         outputs=output_text,
         queue=False,
@@ -153,6 +181,8 @@ with gr.Blocks() as demo:
             temperature,
             top_p,
             top_k,
+            no_repeat_ngram_size,
+            num_beams,
         ],
         outputs=[input_text, output_text],
         queue=False,