Merge pull request #124 from 00x808080/feature/add-model-constants-file

Add model constants file and standardize model name usage

Author: OKUA1

skllm/llm/anthropic/completion.py

@@ -1,12 +1,13 @@
 from typing import Dict, List, Optional
 from skllm.llm.anthropic.credentials import set_credentials
 from skllm.utils import retry
+from model_constants import ANTHROPIC_CLAUDE_MODEL
 
 @retry(max_retries=3)
 def get_chat_completion(
     messages: List[Dict],
     key: str,
-    model: str = "claude-3-haiku-20240307",
+    model: str = ANTHROPIC_CLAUDE_MODEL,
     max_tokens: int = 1000,
     temperature: float = 0.0,
     system: Optional[str] = None,

skllm/llm/gpt/clients/openai/completion.py

@@ -5,14 +5,15 @@
     set_credentials,
 )
 from skllm.utils import retry
+from model_constants import OPENAI_GPT_MODEL
 
 
 @retry(max_retries=3)
 def get_chat_completion(
     messages: dict,
     key: str,
     org: str,
-    model: str = "gpt-3.5-turbo",
+    model: str = OPENAI_GPT_MODEL,
     api="openai",
     json_response=False,
 ):
@@ -27,7 +28,7 @@ def get_chat_completion(
     org : str
         The OPEN AI organization ID to use.
     model : str, optional
-        The OPEN AI model to use. Defaults to "gpt-3.5-turbo".
+        The OPEN AI model to use.
     max_retries : int, optional
         The maximum number of retries to use. Defaults to 3.
     api : str

skllm/llm/gpt/clients/openai/embedding.py

@@ -2,14 +2,15 @@
 from skllm.utils import retry
 import openai
 from openai import OpenAI
+from model_constants import OPENAI_EMBEDDING_MODEL
 
 
 @retry(max_retries=3)
 def get_embedding(
     text: str,
     key: str,
     org: str,
-    model: str = "text-embedding-ada-002",
+    model: str = OPENAI_EMBEDDING_MODEL,
     api: str = "openai"
 ):
     """
@@ -24,7 +25,7 @@ def get_embedding(
     org : str
         The OPEN AI organization ID to use.
     model : str, optional
-        The model to use. Defaults to "text-embedding-ada-002".
+        The model to use.
     max_retries : int, optional
         The maximum number of retries to use. Defaults to 3.
     api: str, optional

skllm/llm/gpt/completion.py

@@ -7,13 +7,14 @@
 )
 from skllm.llm.gpt.utils import split_to_api_and_model
 from skllm.config import SKLLMConfig as _Config
+from model_constants import OPENAI_GPT_MODEL
 
 
 def get_chat_completion(
     messages: dict,
     openai_key: str = None,
     openai_org: str = None,
-    model: str = "gpt-3.5-turbo",
+    model: str = OPENAI_GPT_MODEL,
     json_response: bool = False,
 ):
     """Gets a chat completion from the OpenAI compatible API."""

skllm/llm/gpt/embedding.py

@@ -1,11 +1,13 @@
 from skllm.llm.gpt.clients.openai.embedding import get_embedding as _oai_get_embedding
 from skllm.llm.gpt.utils import split_to_api_and_model
+from model_constants import OPENAI_EMBEDDING_MODEL
+
 
 def get_embedding(
     text: str,
     key: str,
     org: str,
-    model: str = "text-embedding-ada-002",
+    model: str = OPENAI_EMBEDDING_MODEL,
 ):
     """
     Encodes a string and return the embedding for a string.
@@ -19,7 +21,7 @@ def get_embedding(
     org : str
         The OPEN AI organization ID to use.
     model : str, optional
-        The model to use. Defaults to "text-embedding-ada-002".
+        The model to use.
 
     Returns
     -------

skllm/llm/gpt/mixin.py

@@ -21,6 +21,7 @@
 import numpy as np
 from tqdm import tqdm
 import json
+from model_constants import OPENAI_GPT_TUNABLE_MODEL
 
 
 def construct_message(role: str, content: str) -> dict:
@@ -219,24 +220,16 @@ def _get_embeddings(self, text: np.ndarray) -> List[List[float]]:
 
 # for now this works only with OpenAI
 class GPTTunableMixin(BaseTunableMixin):
-    _supported_tunable_models = [
-        "gpt-3.5-turbo-0125",
-        "gpt-3.5-turbo",
-        "gpt-4o-mini-2024-07-18",
-        "gpt-4o-mini",
-    ]
-
     def _build_label(self, label: str):
         return json.dumps({"label": label})
 
     def _set_hyperparameters(self, base_model: str, n_epochs: int, custom_suffix: str):
-        self.base_model = base_model
+        self.base_model = OPENAI_GPT_TUNABLE_MODEL
         self.n_epochs = n_epochs
         self.custom_suffix = custom_suffix
         if base_model not in self._supported_tunable_models:
             raise ValueError(
                 f"Model {base_model} is not supported. Supported models are"
-                f" {self._supported_tunable_models}"
             )
 
     def _tune(self, X, y):

skllm/model_constants.py

@@ -0,0 +1,10 @@
+# OpenAI models
+OPENAI_GPT_MODEL = "gpt-3.5-turbo"
+OPENAI_GPT_TUNABLE_MODEL = "gpt-3.5-turbo-0613"
+OPENAI_EMBEDDING_MODEL = "text-embedding-ada-002"
+
+# Anthropic (Claude) models
+ANTHROPIC_CLAUDE_MODEL = "claude-3-haiku-20240307"
+
+# Vertex AI models
+VERTEX_DEFAULT_MODEL = "text-bison@002"      

skllm/models/anthropic/classification/few_shot.py

@@ -9,14 +9,15 @@
 from skllm.models._base.vectorizer import BaseVectorizer
 from skllm.memory.base import IndexConstructor
 from typing import Optional
+from model_constants import ANTHROPIC_CLAUDE_MODEL, OPENAI_EMBEDDING_MODEL
 
 
 class FewShotClaudeClassifier(BaseFewShotClassifier, ClaudeClassifierMixin, SingleLabelMixin):
     """Few-shot text classifier using Anthropic's Claude API for single-label classification tasks."""
 
     def __init__(
         self,
-        model: str = "claude-3-haiku-20240307",
+        model: str = ANTHROPIC_CLAUDE_MODEL,
         default_label: str = "Random",
         prompt_template: Optional[str] = None,
         key: Optional[str] = None,
@@ -28,7 +29,7 @@ def __init__(
         Parameters
         ----------
         model : str, optional
-            model to use, by default "claude-3-haiku-20240307"
+            model to use
         default_label : str, optional
             default label for failed prediction; if "Random" -> selects randomly based on class frequencies
         prompt_template : Optional[str], optional
@@ -52,7 +53,7 @@ class MultiLabelFewShotClaudeClassifier(
 
     def __init__(
         self,
-        model: str = "claude-3-haiku-20240307",
+        model: str = ANTHROPIC_CLAUDE_MODEL,
         default_label: str = "Random",
         max_labels: Optional[int] = 5,
         prompt_template: Optional[str] = None,
@@ -65,7 +66,7 @@ def __init__(
         Parameters
         ----------
         model : str, optional
-            model to use, by default "claude-3-haiku-20240307"
+            model to use
         default_label : str, optional
             default label for failed prediction; if "Random" -> selects randomly based on class frequencies
         max_labels : Optional[int], optional
@@ -95,7 +96,7 @@ class DynamicFewShotClaudeClassifier(
 
     def __init__(
         self,
-        model: str = "claude-3-haiku-20240307",
+        model: str = ANTHROPIC_CLAUDE_MODEL,
         default_label: str = "Random",
         prompt_template: Optional[str] = None,
         key: Optional[str] = None,
@@ -112,7 +113,7 @@ def __init__(
         Parameters
         ----------
         model : str, optional
-            model to use, by default "claude-3-haiku-20240307"
+            model to use
         default_label : str, optional
             default label for failed prediction; if "Random" -> selects randomly based on class frequencies
         prompt_template : Optional[str], optional
@@ -129,7 +130,7 @@ def __init__(
             metric used for similarity search, by default "euclidean"
         """
         if vectorizer is None:
-            vectorizer = GPTVectorizer(model="text-embedding-ada-002", key=key)
+            vectorizer = GPTVectorizer(model=OPENAI_EMBEDDING_MODEL, key=key)
         super().__init__(
             model=model,
             default_label=default_label,

skllm/models/anthropic/classification/zero_shot.py

@@ -6,6 +6,7 @@
 )
 from skllm.llm.anthropic.mixin import ClaudeClassifierMixin as _ClaudeClassifierMixin
 from typing import Optional
+from model_constants import ANTHROPIC_CLAUDE_MODEL
 
 
 class ZeroShotClaudeClassifier(
@@ -15,7 +16,7 @@ class ZeroShotClaudeClassifier(
 
     def __init__(
         self,
-        model: str = "claude-3-haiku-20240307",
+        model: str = ANTHROPIC_CLAUDE_MODEL,
         default_label: str = "Random",
         prompt_template: Optional[str] = None,
         key: Optional[str] = None,
@@ -27,7 +28,7 @@ def __init__(
         Parameters
         ----------
         model : str, optional
-            Model to use, by default "claude-3-haiku-20240307".
+            Model to use
         default_label : str, optional
             Default label for failed predictions; if "Random", selects randomly based on class frequencies, defaults to "Random".
         prompt_template : Optional[str], optional
@@ -51,7 +52,7 @@ class CoTClaudeClassifier(
 
     def __init__(
         self,
-        model: str = "claude-3-haiku-20240307",
+        model: str = ANTHROPIC_CLAUDE_MODEL,
         default_label: str = "Random",
         prompt_template: Optional[str] = None,
         key: Optional[str] = None,
@@ -63,7 +64,7 @@ def __init__(
         Parameters
         ----------
         model : str, optional
-            Model to use, by default "claude-3-haiku-20240307".
+            Model to use.
         default_label : str, optional
             Default label for failed predictions; if "Random", selects randomly based on class frequencies, defaults to "Random".
         prompt_template : Optional[str], optional
@@ -87,7 +88,7 @@ class MultiLabelZeroShotClaudeClassifier(
 
     def __init__(
         self,
-        model: str = "claude-3-haiku-20240307",
+        model: str = ANTHROPIC_CLAUDE_MODEL,
         default_label: str = "Random",
         max_labels: Optional[int] = 5,
         prompt_template: Optional[str] = None,
@@ -100,7 +101,7 @@ def __init__(
         Parameters
         ----------
         model : str, optional
-            Model to use, by default "claude-3-haiku-20240307".
+            Model to use.
         default_label : str, optional
             Default label for failed predictions; if "Random", selects randomly based on class frequencies, defaults to "Random".
         max_labels : Optional[int], optional

skllm/models/anthropic/tagging/ner.py

@@ -1,6 +1,7 @@
 from skllm.models._base.tagger import ExplainableNER as _ExplainableNER
 from skllm.llm.anthropic.mixin import ClaudeTextCompletionMixin as _ClaudeTextCompletionMixin
 from typing import Optional, Dict
+from model_constants import ANTHROPIC_CLAUDE_MODEL
 
 
 class AnthropicExplainableNER(_ExplainableNER, _ClaudeTextCompletionMixin):
@@ -11,7 +12,7 @@ def __init__(
         entities: Dict[str, str],
         display_predictions: bool = False,
         sparse_output: bool = True,
-        model: str = "claude-3-haiku-20240307",
+        model: str = ANTHROPIC_CLAUDE_MODEL,
         key: Optional[str] = None,
         num_workers: int = 1,
     ) -> None:
@@ -27,7 +28,7 @@ def __init__(
         sparse_output : bool, optional
             whether to generate a sparse representation of the predictions, by default True
         model : str, optional
-            model to use, by default "claude-3-haiku-20240307"
+            model to use
         key : Optional[str], optional
             estimator-specific API key; if None, retrieved from the global config
         num_workers : int, optional