⚡️ adjust params for csearch

Signed-off-by: peter szemraj <[email protected]>

app.py

@@ -20,31 +20,28 @@ def generate_text(
     gen_length=64,
     penalty_alpha=0.6,
     top_k=6,
-    no_repeat_ngram_size=2,
     length_penalty=1.0,
     # perma params (not set by user)
     abs_max_length=512,
     verbose=False,
 ):
     """
-    generate_text - generate text from a prompt using a text generation pipeline
-
-    Args:
-        prompt (str): the prompt to generate text from
-        model_input (_type_): the text generation pipeline
-        max_length (int, optional): the maximum length of the generated text. Defaults to 128.
-        method (str, optional): the generation method. Defaults to "Sampling".
-        verbose (bool, optional): the verbosity of the output. Defaults to False.
-
-    Returns:
-        str: the generated text
+    generate_text - generate text using the text generation pipeline
+
+    :param str prompt: the prompt to use for the text generation pipeline
+    :param int gen_length: the number of tokens to generate
+    :param float penalty_alpha: the penalty alpha for the text generation pipeline (contrastive search)
+    :param int top_k: the top k for the text generation pipeline (contrastive search)
+    :param int abs_max_length: the absolute max length for the text generation pipeline
+    :param bool verbose: verbose output
+    :return str: the generated text
     """
     global generator
     if verbose:
         logging.info(f"Generating text from prompt:\n\n{prompt}")
         logging.info(
             pp.pformat(
-                f"params:\tmax_length={gen_length}, penalty_alpha={penalty_alpha}, top_k={top_k}, no_repeat_ngram_size={no_repeat_ngram_size}, length_penalty={length_penalty}"
+                f"params:\tmax_length={gen_length}, penalty_alpha={penalty_alpha}, top_k={top_k}, length_penalty={length_penalty}"
             )
         )
     st = time.perf_counter()
@@ -59,7 +56,6 @@ def generate_text(
         min_length=input_len + 4,
         penalty_alpha=penalty_alpha,
         top_k=top_k,
-        no_repeat_ngram_size=no_repeat_ngram_size,
         length_penalty=length_penalty,
     )  # generate
     response = result[0]["generated_text"]
@@ -107,15 +103,14 @@ def get_parser():
         default="postbot/distilgpt2-emailgen-V2",
         help="Pass an different huggingface model tag to use a custom model",
     )
-
     parser.add_argument(
-        "-v",
-        "--verbose",
+        "-l",
+        "--max_length",
         required=False,
-        action="store_true",
-        help="Verbose output",
+        type=int,
+        default=64,
+        help="default max length of the generated text",
     )
-
     parser.add_argument(
         "-a",
         "--penalty_alpha",
@@ -131,6 +126,13 @@ def get_parser():
         default=6,
         help="The top k for the text generation pipeline (contrastive search) - default 6",
     )
+    parser.add_argument(
+        "-v",
+        "--verbose",
+        required=False,
+        action="store_true",
+        help="Verbose output",
+    )
     return parser
 
 
@@ -192,7 +194,7 @@ if __name__ == "__main__":
                 )
                 num_gen_tokens = gr.Slider(
                     label="Generation Tokens",
-                    value=32,
+                    value=40,
                     maximum=96,
                     minimum=16,
                     step=8,
@@ -217,7 +219,7 @@ if __name__ == "__main__":
             gr.Markdown("---")
             gr.Markdown("## Advanced Options")
             gr.Markdown(
-                "This demo generates text via the new [constrastive search](https://huggingface.co/blog/introducing-csearch). See details on the csearch blog post for the methods' parameters or [here](https://huggingface.co/blog/how-to-generate), for general decoding."
+                "This demo generates text via the new [contrastive search](https://huggingface.co/blog/introducing-csearch). See the csearch blog post for details on the parameters or [here](https://huggingface.co/blog/how-to-generate), for general decoding."
             )
             with gr.Row():
                 model_name = gr.Dropdown(
@@ -229,11 +231,6 @@ if __name__ == "__main__":
                     "Load Model",
                     variant="secondary",
                 )
-                no_repeat_ngram_size = gr.Radio(
-                    choices=[1, 2, 3, 4],
-                    label="no repeat ngram size",
-                    value=3,
-                )
             with gr.Row():
                 contrastive_top_k = gr.Radio(
                     choices=[2, 4, 6, 8],
@@ -280,7 +277,6 @@ if __name__ == "__main__":
                 num_gen_tokens,
                 penalty_alpha,
                 contrastive_top_k,
-                no_repeat_ngram_size,
                 length_penalty,
             ],
             outputs=[email_mailto_button, generated_email],