Merge pull request #573 from GooeyAI/max_tokens_and_max_completion_to…

…kens use max_completion_tokens for o1 models, and max_tokens for the rest
GooeyAI · Dec 24, 2024 · 9668df9 · 9668df9
2 parents 6d0637d + b36524f
commit 9668df9
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/daras_ai_v2/language_model.py b/daras_ai_v2/language_model.py
@@ -1049,6 +1049,10 @@ def run_openai_chat(
         for entry in messages:
             if entry["role"] == CHATML_ROLE_SYSTEM:
                 entry["role"] = CHATML_ROLE_USER
+        max_tokens = NOT_GIVEN
+    else:
+        max_tokens = max_completion_tokens
+        max_completion_tokens = NOT_GIVEN
 
     if avoid_repetition:
         frequency_penalty = 0.1
@@ -1063,6 +1067,7 @@ def run_openai_chat(
             _get_chat_completions_create(
                 model=model_str,
                 messages=messages,
+                max_tokens=max_tokens,
                 max_completion_tokens=max_completion_tokens,
                 stop=stop or NOT_GIVEN,
                 n=num_outputs,