[TLM] Release o1 models (#353)

cleanlab · Jan 29, 2025 · 6be1bcb · 6be1bcb
1 parent 76f479a
commit 6be1bcb
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 1 deletion.
diff --git a/cleanlab_studio/internal/constants.py b/cleanlab_studio/internal/constants.py
@@ -9,6 +9,8 @@
     "gpt-4o",
     "gpt-4o-mini",
     "o1-preview",
+    "o1",
+    "o1-mini",
     "claude-3-haiku",
     "claude-3.5-haiku",
     "claude-3-sonnet",

diff --git a/cleanlab_studio/studio/trustworthy_language_model.py b/cleanlab_studio/studio/trustworthy_language_model.py
@@ -795,7 +795,7 @@ class TLMOptions(TypedDict):
     Args:
         model (str, default = "gpt-4o-mini"): underlying base LLM to use (better models yield better results, faster models yield faster/cheaper results).
         - Models currently supported include: "gpt-4o-mini", "gpt-4o", "o1-preview", "gpt-3.5-turbo-16k", "gpt-4", "claude-3.5-sonnet", "claude-3-haiku".
-        - Additional models supported in beta include: "claude-3.5-sonnet-v2", "claude-3.5-haiku", "nova-micro", "nova-lite", "nova-pro".
+        - Additional models supported in beta include: "o1", "o1-mini", "claude-3.5-sonnet-v2", "claude-3.5-haiku", "nova-micro", "nova-lite", "nova-pro".
 
         max_tokens (int, default = 512): the maximum number of tokens to generate in the TLM response.
         This number will impact the maximum number of tokens you will see in the output response, and also the number of tokens

diff --git a/tests/tlm/constants.py b/tests/tlm/constants.py
@@ -22,6 +22,8 @@
     "claude-3.5-sonnet",
     "claude-3.5-sonnet-v2",
     "o1-preview",
+    "o1",
+    "o1-mini",
     "nova-lite",
     "nova-pro",
 ]
@@ -35,6 +37,8 @@
     "claude-3.5-sonnet",
     "claude-3.5-sonnet-v2",
     "o1-preview",
+    "o1",
+    "o1-mini",
     "nova-micro",
     "nova-lite",
     "nova-pro",