✨ add beam groups

Signed-off-by: Peter <[email protected]>

app.py

@@ -23,6 +23,7 @@ def generate_text(
     # perma params (not set by user)
     repetition_penalty=3.5,
     abs_max_length=512,
+    num_beam_groups=2,
     verbose=False,
 ):
     """
@@ -55,6 +56,7 @@ def generate_text(
         max_length=gen_length + input_len,
         min_length=input_len + 4,
         num_beams=num_beams,
+        num_beam_groups=num_beam_groups,
         repetition_penalty=repetition_penalty,
         no_repeat_ngram_size=no_repeat_ngram_size,
         length_penalty=length_penalty,
@@ -111,6 +113,21 @@ def get_parser():
         action="store_true",
         help="Verbose output",
     )
+
+    parser.add_argument(
+    "-nb",
+        "--num_beams",
+        type=int,
+        default=4,
+        help="Number of beams for beam search. 1 means no beam search.",
+    )
+
+    parser.add_argument(
+        "--num_beam_groups",
+        type=int,
+        default=2,
+        help="Number of groups to divide nbest candidates into in order to ensure diversity among different groups of beams that yield the best n results. 1 means no group beam search.",
+    )
     return parser