From 3849856c2ec0fe09ab2a0cfd2be75190ab650faf Mon Sep 17 00:00:00 2001
From: Hui Wen <45724323+huiwengoh@users.noreply.github.com>
Date: Fri, 10 Jan 2025 17:26:25 -0500
Subject: [PATCH] [TLM] Add more bedrock models (#344)

---
 cleanlab_studio/internal/constants.py      | 11 ++++++-----
 cleanlab_studio/internal/tlm/validation.py |  5 +++--
 cleanlab_studio/version.py                 |  2 +-
 tests/tlm/conftest.py                      |  4 +++-
 tests/tlm/constants.py                     | 14 +++++++++++++-
 5 files changed, 26 insertions(+), 10 deletions(-)

diff --git a/cleanlab_studio/internal/constants.py b/cleanlab_studio/internal/constants.py
index 385d93cb..4655a7cf 100644
--- a/cleanlab_studio/internal/constants.py
+++ b/cleanlab_studio/internal/constants.py
@@ -10,17 +10,18 @@
     "gpt-4o-mini",
     "o1-preview",
     "claude-3-haiku",
+    "claude-3.5-haiku",
     "claude-3-sonnet",
     "claude-3.5-sonnet",
+    "claude-3.5-sonnet-v2",
+    "nova-micro",
+    "nova-lite",
+    "nova-pro",
 ]
 _TLM_DEFAULT_MODEL: str = "gpt-4o-mini"
 _TLM_MAX_RETRIES: int = 3  # TODO: finalize this number
 _TLM_MAX_TOKEN_RANGE: Dict[str, Tuple[int, int]] = {  # model: (min, max)
-    "gpt-3.5-turbo-16k": (64, 4096),
-    "gpt-4": (64, 4096),
-    "gpt-4o": (64, 4096),
-    "gpt-4o-mini": (64, 4096),
-    "o1-preview": (64, 4096),
+    "default": (64, 4096),
     "claude-3-haiku": (64, 512),
     "claude-3-sonnet": (64, 512),
     "claude-3.5-sonnet": (64, 512),
diff --git a/cleanlab_studio/internal/tlm/validation.py b/cleanlab_studio/internal/tlm/validation.py
index 3f012367..69468368 100644
--- a/cleanlab_studio/internal/tlm/validation.py
+++ b/cleanlab_studio/internal/tlm/validation.py
@@ -173,9 +173,10 @@ def validate_tlm_options(options: Any) -> None:
                 raise ValidationError(f"Invalid type {type(val)}, max_tokens must be an integer")
 
             model = options.get("model", _TLM_DEFAULT_MODEL)
-            if val < _TLM_MAX_TOKEN_RANGE[model][0] or val > _TLM_MAX_TOKEN_RANGE[model][1]:
+            max_tokens_range = _TLM_MAX_TOKEN_RANGE.get(model, _TLM_MAX_TOKEN_RANGE["default"])
+            if val < max_tokens_range[0] or val > max_tokens_range[1]:
                 raise ValidationError(
-                    f"Invalid value {val}, max_tokens for {model} must be in the range {_TLM_MAX_TOKEN_RANGE[model]}"
+                    f"Invalid value {val}, max_tokens for {model} must be in the range {max_tokens_range}"
                 )
 
         elif option == "model":
diff --git a/cleanlab_studio/version.py b/cleanlab_studio/version.py
index d7ca3e3d..dc76603b 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.11"
+__version__ = "2.5.12"
 
 SCHEMA_VERSION = "0.2.0"
 MIN_SCHEMA_VERSION = "0.1.0"
diff --git a/tests/tlm/conftest.py b/tests/tlm/conftest.py
index a0a514e7..2fa8af17 100644
--- a/tests/tlm/conftest.py
+++ b/tests/tlm/conftest.py
@@ -89,7 +89,9 @@ def _get_options_dictionary(model: Optional[str]) -> dict:
     add_log_perplexity_score = np.random.choice([True, False])
 
     if add_max_tokens:
-        max_tokens = _TLM_MAX_TOKEN_RANGE[options.get("model", _TLM_DEFAULT_MODEL)][1]
+        max_tokens = _TLM_MAX_TOKEN_RANGE.get(
+            options.get("model", _TLM_DEFAULT_MODEL), _TLM_MAX_TOKEN_RANGE["default"]
+        )[1]
         options["max_tokens"] = int(np.random.randint(64, max_tokens))
     if add_use_self_reflection:
         options["use_self_reflection"] = random.choice([True, False])
diff --git a/tests/tlm/constants.py b/tests/tlm/constants.py
index f597b509..3ae754fa 100644
--- a/tests/tlm/constants.py
+++ b/tests/tlm/constants.py
@@ -16,13 +16,25 @@
 CHARACTERS_PER_TOKEN: int = 4
 
 # Property tests for TLM
-excluded_tlm_models: List[str] = ["claude-3-sonnet", "claude-3.5-sonnet", "o1-preview"]
+excluded_tlm_models: List[str] = [
+    "claude-3-sonnet",
+    "claude-3.5-sonnet",
+    "claude-3.5-sonnet-v2",
+    "o1-preview",
+    "nova-lite",
+    "nova-pro",
+]
 VALID_TLM_MODELS: List[str] = [
     model for model in _VALID_TLM_MODELS if model not in excluded_tlm_models
 ]
 MODELS_WITH_NO_PERPLEXITY_SCORE: List[str] = [
     "claude-3-haiku",
+    "claude-3.5-haiku",
     "claude-3-sonnet",
     "claude-3.5-sonnet",
+    "claude-3.5-sonnet-v2",
     "o1-preview",
+    "nova-micro",
+    "nova-lite",
+    "nova-pro",
 ]