feat: add Gemini-pro-1.5 to GeminiTextGenerator Tuning and Support score() method in Gemini-pro-1.5

bigframes/ml/llm.py

@@ -874,7 +874,7 @@ def fit(
         X: utils.ArrayType,
         y: utils.ArrayType,
     ) -> GeminiTextGenerator:
-        """Fine tune GeminiTextGenerator model. Only support "gemini-pro" model for now.
+        """Fine tune GeminiTextGenerator model. Only support "gemini-pro" and "gemini-1.5" models for now.
 
         .. note::
 
@@ -892,8 +892,12 @@ def fit(
         Returns:
             GeminiTextGenerator: Fitted estimator.
         """
-        if self._bqml_model.model_name.startswith("gemini-1.5"):
-            raise NotImplementedError("Fit is not supported for gemini-1.5 model.")
+        # Support gemini-1.5 and gemini-pro
+        supported_models = ["gemini-pro", "gemini-1.5-pro-002", "gemini-1.5-flash-002"]
+        if self.model_name not in supported_models:
+            raise NotImplementedError(
+                "Score is not supported models other than gemini-pro or gemini-1.5 model."
+            )
 
         X, y = utils.batch_convert_to_dataframe(X, y)
 
@@ -1009,7 +1013,7 @@ def score(
             "text_generation", "classification", "summarization", "question_answering"
         ] = "text_generation",
     ) -> bpd.DataFrame:
-        """Calculate evaluation metrics of the model. Only "gemini-pro" model is supported for now.
+        """Calculate evaluation metrics of the model. Only "gemini-pro" and "gemini-1.5" models are supported for now.
 
         .. note::
 
@@ -1041,9 +1045,12 @@ def score(
         if not self._bqml_model:
             raise RuntimeError("A model must be fitted before score")
 
-        # TODO(ashleyxu): Support gemini-1.5 when the rollout is ready. b/344891364.
-        if self._bqml_model.model_name.startswith("gemini-1.5"):
-            raise NotImplementedError("Score is not supported for gemini-1.5 model.")
+        # Support gemini-1.5 and gemini-pro
+        supported_models = ["gemini-pro", "gemini-1.5-pro-002", "gemini-1.5-flash-002"]
+        if self.model_name not in supported_models:
+            raise NotImplementedError(
+                "Score is not supported models other than gemini-pro or gemini-1.5 model."
+            )
 
         X, y = utils.batch_convert_to_dataframe(X, y, session=self._bqml_model.session)
 

tests/system/load/test_llm.py

@@ -38,12 +38,19 @@ def llm_remote_text_df(session, llm_remote_text_pandas_df):
     return session.read_pandas(llm_remote_text_pandas_df)
 
 
-@pytest.mark.flaky(retries=2)
+@pytest.mark.parametrize(
+    "model_name",
+    (
+        "gemini-pro",
+        "gemini-1.5-pro-002",
+        "gemini-1.5-flash-002",
+    ),
+)
 def test_llm_gemini_configure_fit(
-    session, llm_fine_tune_df_default_index, llm_remote_text_df
+    session, model_name, llm_fine_tune_df_default_index, llm_remote_text_df
 ):
     model = llm.GeminiTextGenerator(
-        session=session, model_name="gemini-pro", max_iterations=1
+        session=session, model_name=model_name, max_iterations=1
     )
 
     X_train = llm_fine_tune_df_default_index[["prompt"]]
@@ -69,7 +76,6 @@ def test_llm_gemini_configure_fit(
         ],
         index=3,
     )
-    # TODO(ashleyxu b/335492787): After bqml rolled out version control: save, load, check parameters to ensure configuration was kept
 
 
 @pytest.mark.flaky(retries=2)

tests/system/small/ml/test_llm.py

@@ -413,9 +413,17 @@ def test_llm_palm_score_params(llm_fine_tune_df_default_index):
     )
 
 
-@pytest.mark.flaky(retries=2)
-def test_llm_gemini_pro_score(llm_fine_tune_df_default_index):
-    model = llm.GeminiTextGenerator(model_name="gemini-pro")
+# test score() function for "gemini-pro" and "gemini-1.5" model
+@pytest.mark.parametrize(
+    "model_name",
+    (
+        "gemini-pro",
+        "gemini-1.5-pro-002",
+        "gemini-1.5-flash-002",
+    ),
+)
+def test_llm_gemini_score(llm_fine_tune_df_default_index, model_name):
+    model = llm.GeminiTextGenerator(model_name=model_name)
 
     # Check score to ensure the model was fitted
     score_result = model.score(