format

app.py

@@ -11,6 +11,7 @@ logging.basicConfig(
 
 use_gpu = torch.cuda.is_available()
 
+
 def generate_text(
     prompt: str,
     gen_length=64,
@@ -40,7 +41,7 @@ def generate_text(
     st = time.perf_counter()
 
     input_tokens = generator.tokenizer(prompt)
-    input_len = len(input_tokens['input_ids'])
+    input_len = len(input_tokens["input_ids"])
     if input_len > abs_max_length:
         logging.info(f"Input too long {input_len} > {abs_max_length}, may cause errors")
     result = generator(
@@ -55,9 +56,8 @@ def generate_text(
         early_stopping=True,
         # tokenizer
         truncation=True,
-
-    ) # generate
-    response = result[0]['generated_text']
+    )  # generate
+    response = result[0]["generated_text"]
     rt = time.perf_counter() - st
     if verbose:
         logging.info(f"Generated text: {response}")
@@ -74,12 +74,12 @@ def get_parser():
     )
 
     parser.add_argument(
-        '-m',
-        '--model',
+        "-m",
+        "--model",
         required=False,
         type=str,
         default="postbot/distilgpt2-emailgen",
-        help='Pass an different huggingface model tag to use a custom model',
+        help="Pass an different huggingface model tag to use a custom model",
     )
 
     parser.add_argument(
@@ -91,6 +91,7 @@ def get_parser():
     )
     return parser
 
+
 default_prompt = """
 Hello,
 
@@ -109,7 +110,6 @@ if __name__ == "__main__":
         device=0 if use_gpu else -1,
     )
 
-
     demo = gr.Blocks()
 
     logging.info("launching interface...")
@@ -119,7 +119,9 @@ if __name__ == "__main__":
         gr.Markdown(
             "Enter part of an email, and the model will autocomplete it for you!"
         )
-        gr.Markdown('The model used is [postbot/distilgpt2-emailgen](https://huggingface.co/postbot/distilgpt2-emailgen)')
+        gr.Markdown(
+            "The model used is [postbot/distilgpt2-emailgen](https://huggingface.co/postbot/distilgpt2-emailgen)"
+        )
         gr.Markdown("---")
 
         with gr.Column():
@@ -151,10 +153,11 @@ if __name__ == "__main__":
                 value=2,
             )
             length_penalty = gr.Slider(
-                            minimum=0.5, maximum=1.0, label="length penalty", default=0.8, step=0.05
+                minimum=0.5, maximum=1.0, label="length penalty", default=0.8, step=0.05
             )
             generated_email = gr.Textbox(
-                label="Generated Result", placeholder="The completed email will appear here"
+                label="Generated Result",
+                placeholder="The completed email will appear here",
             )
 
             generate_button = gr.Button(
@@ -168,16 +171,24 @@ if __name__ == "__main__":
             gr.Markdown(
                 "This model is a fine-tuned version of distilgpt2 on a dataset of 50k emails sourced from the internet, including the classic `aeslc` dataset."
             )
-            gr.Markdown("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; 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(
+                "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; 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("---")
 
         generate_button.click(
             fn=generate_text,
-            inputs=[prompt_text, num_gen_tokens, num_beams, no_repeat_ngram_size, length_penalty],
+            inputs=[
+                prompt_text,
+                num_gen_tokens,
+                num_beams,
+                no_repeat_ngram_size,
+                length_penalty,
+            ],
             outputs=[generated_email],
         )
 
     demo.launch(
         enable_queue=True,
-        share=True, # for local testing
+        share=True,  # for local testing
     )