generate with contrastive search

Signed-off-by: peter szemraj <peterszemraj@gmail.com>

app.py

@@ -17,12 +17,11 @@ use_gpu = torch.cuda.is_available()
 def generate_text(
     prompt: str,
     gen_length=64,
-    num_beams=4,
+    penalty_alpha=0.6,
+    top_k=6,
     no_repeat_ngram_size=2,
     length_penalty=1.0,
-    num_beam_groups=1,
     # perma params (not set by user)
-    repetition_penalty=3.5,
     abs_max_length=512,
     verbose=False,
 ):
@@ -53,15 +52,13 @@ def generate_text(
         logging.info(f"Input too long {input_len} > {abs_max_length}, may cause errors")
     result = generator(
         prompt,
-        max_length=gen_length + input_len,
+        max_new_tokens=gen_length,
+        max_length=None,  # in case of default max_length
         min_length=input_len + 4,
-        num_beams=num_beams,
-        num_beam_groups=num_beam_groups,
-        repetition_penalty=repetition_penalty,
+        penalty_alpha=penalty_alpha,
+        top_k=top_k,
         no_repeat_ngram_size=no_repeat_ngram_size,
         length_penalty=length_penalty,
-        do_sample=False,
-        early_stopping=True,
     )  # generate
     response = result[0]["generated_text"]
     rt = time.perf_counter() - st
@@ -118,18 +115,19 @@ def get_parser():
     )
 
     parser.add_argument(
-        "-nb",
-        "--num_beams",
-        type=int,
-        default=4,
-        help="Number of beams for beam search. 1 means no beam search.",
+        "-a",
+        "--penalty_alpha",
+        type=float,
+        default=0.6,
+        help="The penalty alpha for the text generation pipeline (contrastive search) - default 0.6",
     )
 
     parser.add_argument(
-        "--num_beam_groups",
+        "-k",
+        "--top_k",
         type=int,
-        default=1,
-        help="Number of groups to divide best candidates into in order to ensure diversity among different groups of beams that yield the best n results. 1 means no group beam search. (default 1)",
+        default=6,
+        help="The top k for the text generation pipeline (contrastive search) - default 6",
     )
     return parser
 
@@ -146,11 +144,18 @@ available_models = [
 ]
 
 if __name__ == "__main__":
+
     logging.info("\n\n\nStarting new instance of app.py")
     args = get_parser().parse_args()
     logging.info(f"received args:\t{args}")
     model_tag = args.model
     verbose = args.verbose
+    top_k = args.top_k
+    alpha = args.penalty_alpha
+
+    assert top_k > 0, "top_k must be greater than 0"
+    assert alpha >= 0.0 and alpha <= 1.0, "penalty_alpha must be between 0 and 1"
+
     logging.info(f"Loading model: {model_tag}, use GPU = {use_gpu}")
     generator = pipeline(
         "text-generation",
@@ -228,16 +233,18 @@ if __name__ == "__main__":
                     value=2,
                 )
             with gr.Row():
-                num_beams = gr.Radio(
-                    choices=[2, 4, 8],
+                contrastive_top_k = gr.Radio(
+                    choices=[2, 4, 6, 8],
                     label="Number of Beams",
-                    value=4,
+                    value=top_k,
                 )
 
-                num_beam_groups = gr.Radio(
-                    choices=[1, 2],
-                    label="Number of Beam Groups",
-                    value=1,
+                penalty_alpha = gr.Slider(
+                    label="Penalty Alpha",
+                    value=alpha,
+                    maximum=1.0,
+                    minimum=0.0,
+                    step=0.1,
                 )
                 length_penalty = gr.Slider(
                     minimum=0.5,
@@ -269,10 +276,10 @@ if __name__ == "__main__":
             inputs=[
                 prompt_text,
                 num_gen_tokens,
-                num_beams,
+                penalty_alpha,
+                contrastive_top_k,
                 no_repeat_ngram_size,
                 length_penalty,
-                num_beam_groups,
             ],
             outputs=[email_mailto_button, generated_email],
         )

requirements.txt

@@ -1,3 +1,3 @@
 gradio
 torch
-transformers
+transformers>=4.24.0