Merge pull request #108 from iryna-kondr/gpt4omini_tuning

Added gpt-4o-mini to the list of tunable models

pyproject.toml

@@ -11,7 +11,7 @@ dependencies = [
   "google-cloud-aiplatform[pipelines]>=1.27.0,<2.0.0"
 ]
 name = "scikit-llm"
-version = "1.3.0"
+version = "1.3.1"
 authors = [
   { name="Oleh Kostromin", email="[email protected]" },
   { name="Iryna Kondrashchenko", email="[email protected]" },

skllm/__init__.py

@@ -1,2 +1,2 @@
-__version__ = '1.3.0'
+__version__ = '1.3.1'
 __author__ = 'Iryna Kondrashchenko, Oleh Kostromin'

skllm/llm/gpt/mixin.py

@@ -66,10 +66,9 @@ def _set_keys(self, key: Optional[str] = None, org: Optional[str] = None) -> Non
         """
         Set the OpenAI key and organization.
         """
-            
+
         self.key = key
         self.org = org
-
 
     def _get_openai_key(self) -> str:
         """
@@ -214,7 +213,12 @@ 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"]
+    _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})
@@ -264,4 +268,4 @@ def _tune(self, X, y):
         self.openai_model = job.fine_tuned_model
         self.model = self.openai_model  # openai_model is probably not required anymore
         delete_file(client, job.training_file)
-        print(f"Finished training.")
+        print(f"Finished training.")