fix: openai requisition

iarahealth · Jun 17, 2024 · e90d120 · e90d120
1 parent 3302275
commit e90d120
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 16 deletions.
diff --git a/run_txt_gen.py b/run_txt_gen.py
@@ -47,7 +47,7 @@
     parser.add_argument(
         "--model",
         type=str,
-        default="gpt-3.5-turbo-16k",
+        default="gpt-3.5-turbo-0125",
         help="ChatGPT model to use",
     )
     parser.add_argument(
@@ -68,7 +68,7 @@
 
     if args.query is None:
         if args.return_type == "frases":
-            args.query = f"Você é um médico laudando. No contexto de {args.context}, gere {args.num} {args.return_type} contendo o termo '[MASK]', separadas por nova linha."
+            args.query = f"Você é um médico ditando o laudo de um paciente. No contexto de {args.context}, gere {args.num} {args.return_type} contendo o termo '[MASK]', separadas por nova linha."
         else:
             args.query = f"No contexto de {args.context}, gere {args.num} {args.return_type} separadas por nova linha."
     else:

diff --git a/text/gen.py b/text/gen.py
@@ -11,24 +11,21 @@
 MAX_TOKENS = {
     # https://platform.openai.com/docs/models/gpt-4
     # https://platform.openai.com/docs/models/gpt-3-5
-    "gpt-4": 8192,
-    "gpt-4-32k": 32768,
-    "gpt-3.5-turbo": 4097,
-    "gpt-3.5-turbo-16k": 16385,
+    "gpt-4-turbo": 128000,
+    "gpt-3.5-turbo-0125": 16385,
 }
 
 
 # https://github.com/openai/openai-cookbook/blob/main/examples/How_to_handle_rate_limits.ipynb
 @retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6))
 def make_chatgpt_query(
-    client, query: str, return_type: str, model: str = "gpt-3.5-turbo"
+    client, query: str, return_type: str, model: str = "gpt-3.5-turbo-0125"
 ) -> List[str]:
     """
     Makes a query to the ChatGPT model and returns the generated response.
 
     Args:
         query (str): The user's query.
-        api_key (str): The API key for accessing the ChatGPT model.
         model (str): The name of the ChatGPT model.
 
     Returns:
@@ -39,11 +36,8 @@ def make_chatgpt_query(
         model=model, messages=[{"role": "user", "content": query}]
     )
     is_truncated = response.usage.total_tokens >= max_tokens
-    response = [
-        line
-        for line in response.choices[0].message.content.split("\n")
-        if line.strip() != ""
-    ]
-    if is_truncated and len(response) > 0:
-        response.pop()
-    return response
+    response_text = response.choices[0].message.content
+    response_lines = [line for line in response_text.split("\n") if line.strip() != ""]
+    if is_truncated and len(response_lines) > 0:
+        response_lines.pop()
+    return response_lines