From e0bf92048736e2f29eaf3bcba00b09376a320267 Mon Sep 17 00:00:00 2001
From: Jay Zhang <jas2600@gmail.com>
Date: Mon, 7 Oct 2024 10:51:15 -0700
Subject: [PATCH] tlm add o1-preview

---
 cleanlab_studio/internal/constants.py | 2 ++
 cleanlab_studio/version.py            | 2 +-
 tests/tlm/constants.py                | 3 ++-
 3 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/cleanlab_studio/internal/constants.py b/cleanlab_studio/internal/constants.py
index 97c4afe7..3a1501e1 100644
--- a/cleanlab_studio/internal/constants.py
+++ b/cleanlab_studio/internal/constants.py
@@ -8,6 +8,7 @@
     "gpt-4",
     "gpt-4o",
     "gpt-4o-mini",
+    "o1-preview",
     "claude-3-haiku",
     "claude-3-sonnet",
     "claude-3.5-sonnet",
@@ -19,6 +20,7 @@
     "gpt-4": (64, 4096),
     "gpt-4o": (64, 4096),
     "gpt-4o-mini": (64, 4096),
+    "o1-preview": (64, 4096),
     "claude-3-haiku": (64, 512),
     "claude-3-sonnet": (64, 512),
     "claude-3.5-sonnet": (64, 512),
diff --git a/cleanlab_studio/version.py b/cleanlab_studio/version.py
index a2da2fc1..48c19ef1 100644
--- a/cleanlab_studio/version.py
+++ b/cleanlab_studio/version.py
@@ -1,7 +1,7 @@
 # Note to developers:
 # Consider if backend's MIN_CLI_VERSION needs updating when pushing any changes to this file.
 
-__version__ = "2.5.4"
+__version__ = "2.5.5"
 
 SCHEMA_VERSION = "0.2.0"
 MIN_SCHEMA_VERSION = "0.1.0"
diff --git a/tests/tlm/constants.py b/tests/tlm/constants.py
index 3ff2fe4a..f597b509 100644
--- a/tests/tlm/constants.py
+++ b/tests/tlm/constants.py
@@ -16,7 +16,7 @@
 CHARACTERS_PER_TOKEN: int = 4
 
 # Property tests for TLM
-excluded_tlm_models: List[str] = ["claude-3-sonnet", "claude-3.5-sonnet"]
+excluded_tlm_models: List[str] = ["claude-3-sonnet", "claude-3.5-sonnet", "o1-preview"]
 VALID_TLM_MODELS: List[str] = [
     model for model in _VALID_TLM_MODELS if model not in excluded_tlm_models
 ]
@@ -24,4 +24,5 @@
     "claude-3-haiku",
     "claude-3-sonnet",
     "claude-3.5-sonnet",
+    "o1-preview",
 ]