Merge branch 'main' of https://huggingface.co/spaces/postbot/autocomplete-emails

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,
@@ -70,7 +72,8 @@ def generate_text(
     formatted_email = postprocess(response)
     return make_mailto_form(body=formatted_email)
 
-def load_emailgen_model(model_tag:str):
+
+def load_emailgen_model(model_tag: str):
     """
     load_emailgen_model - load a text generation pipeline for email generation
 
@@ -87,6 +90,7 @@ def load_emailgen_model(model_tag:str):
         device=0 if use_gpu else -1,
     )
 
+
 def get_parser():
     """
     get_parser - a helper function for the argparse module
@@ -111,6 +115,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
 
 
@@ -119,7 +138,11 @@ Hello,
 
 Following up on last week's bubblegum shipment, I"""
 
-available_models = ['postbot/distilgpt2-emailgen-V2', 'postbot/distilgpt2-emailgen', 'postbot/gpt2-medium-emailgen']
+available_models = [
+    "postbot/distilgpt2-emailgen-V2",
+    "postbot/distilgpt2-emailgen",
+    "postbot/gpt2-medium-emailgen",
+]
 
 if __name__ == "__main__":
     logging.info("\n\n\nStarting new instance of app.py")
@@ -189,16 +212,18 @@ if __name__ == "__main__":
                     value=model_tag,
                 )
                 load_model_button = gr.Button(
-                    'Load Model',
-                    variant='secondary',
+                    "Load Model",
+                    variant="secondary",
                 )
 
             with gr.Row():
                 num_beams = gr.Radio(
-                                    choices=[4, 8, 12, 16],
-                                    label="Number of Beams",
-                                    value=4,
-                                )
+                    choices=[4, 8, 12, 16],
+                    label="Number of Beams",
+                    value=4,
+                )
+            with gr.Row():
+
                 no_repeat_ngram_size = gr.Radio(
                     choices=[1, 2, 3, 4],
                     label="no repeat ngram size",

utils.py

@@ -4,6 +4,7 @@
 import logging
 import re
 
+
 def postprocess(text: str):
     """
     postprocess - remove common values in scraped dataset