:art: format code to black

ai_single_response.py

@@ -122,7 +122,7 @@ def query_gpt_model(
 ):
     """
     query_gpt_model - the main function that calls the model.
-    
+
     Parameters:
     -----------
     prompt_msg (str): the prompt to be sent to the model
@@ -137,7 +137,7 @@ def query_gpt_model(
     use_gpu (bool, optional): use gpu. Defaults to False.
     """
     ai = aitextgen(
-        model="pszemraj/Ballpark-Trivia-L", # THIS WORKS
+        model="pszemraj/Ballpark-Trivia-L",  # THIS WORKS
         # model="pszemraj/Ballpark-Trivia-XL", # does not seem to work
         to_gpu=use_gpu,
     )
@@ -155,22 +155,22 @@ def query_gpt_model(
     # call the model
     print("\n... generating...")
     this_result = ai.generate(
-                                n=1,
-                                top_k=kparam,
-                                batch_size=batch_size,
-                                # the prompt input counts for text length constraints
-                                max_length=resp_length + pr_len,
-                                min_length=resp_min + pr_len,
-                                prompt=this_prompt,
-                                temperature=temp,
-                                top_p=top_p,
-                                top_k=kparam,
-                                do_sample=True,
-                                return_as_list=True,
-                                use_cache=True,
-                                num_beams=beams,
-                                no_repeat_ngram_size=2,
-                            )
+        n=1,
+        top_k=kparam,
+        batch_size=batch_size,
+        # the prompt input counts for text length constraints
+        max_length=resp_length + pr_len,
+        min_length=resp_min + pr_len,
+        prompt=this_prompt,
+        temperature=temp,
+        top_p=top_p,
+        top_k=kparam,
+        do_sample=True,
+        return_as_list=True,
+        use_cache=True,
+        num_beams=beams,
+        no_repeat_ngram_size=2,
+    )
     if verbose:
         print("\n... generated:\n")
         pp.pprint(this_result)  # for debugging