✨ add ability to change models

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

app.py

@@ -72,6 +72,24 @@ def generate_text(
     formatted_email = postprocess(response)
     return formatted_email, make_email_link(body=formatted_email)
 
+def load_emailgen_model(model_tag:str):
+    """
+    load_emailgen_model - load a text generation pipeline for email generation
+
+    Args:
+        model_tag (str): the huggingface model tag to load
+
+    Returns:
+        transformers.pipelines.TextGenerationPipeline: the text generation pipeline
+    """
+
+    generator = pipeline(
+        "text-generation",
+        model_tag,
+        device=0 if use_gpu else -1,
+    )
+
+    return generator
 
 def get_parser():
     """
@@ -105,6 +123,8 @@ Hello,
 
 Following up on last week's bubblegum shipment, I"""
 
+available_models = ['postbot/distilgpt2-emailgen', 'postbot/distilgpt2-emailgen-V2', 'postbot/gpt2-medium-emailgen']
+
 if __name__ == "__main__":
     logging.info("\n\n\nStarting new instance of app.py")
     args = get_parser().parse_args()
@@ -166,12 +186,21 @@ if __name__ == "__main__":
             gr.Markdown(
                 "This demo generates text via beam search. See details about these parameters [here](https://huggingface.co/blog/how-to-generate), otherwise they should be fine as-is."
             )
-
-            num_beams = gr.Radio(
-                choices=[4, 8, 16],
-                label="Number of Beams",
-                value=4,
-            )
+            with gr.Row():
+                model_name = gr.Dropdown(
+                    choices=available_models,
+                    label="Choose a model",
+                    value=model_tag,
+                )
+                load_model_button = gr.Button(
+                    'Load Model',
+                    variant='secondary',
+                )
+                num_beams = gr.Radio(
+                    choices=[4, 8, 16],
+                    label="Number of Beams",
+                    value=4,
+                )
             with gr.Row():
 
                 no_repeat_ngram_size = gr.Radio(
@@ -198,6 +227,7 @@ if __name__ == "__main__":
                 "The intended use of this model is to provide suggestions to _auto-complete_ the rest of your email. Said another way, it should serve as a **tool to write predictable emails faster**. It is not intended to write entire emails from scratch; at least **some input** is required to guide the direction of the model.\n\nPlease verify any suggestions by the model for A) False claims and B) negation statements **before** accepting/sending something."
             )
             gr.Markdown("---")
+
         clear_button.click(
             fn=clear,
             inputs=[prompt_text],
@@ -215,6 +245,11 @@ if __name__ == "__main__":
             outputs=[generated_email, email_link],
         )
 
+        load_model_button.click(
+            fn=load_emailgen_model,
+            inputs=[model_name],
+            outputs=[generator],
+        )
     demo.launch(
         enable_queue=True,
         share=True,  # for local testing