address comments

bigframes/ml/llm.py

@@ -29,7 +29,6 @@
 
 _BQML_PARAMS_MAPPING = {
     "max_iterations": "maxIterations",
-    "evaluation_task": "evaluationTask",
 }
 
 _TEXT_GENERATOR_BISON_ENDPOINT = "text-bison"
@@ -75,9 +74,6 @@ class PaLM2TextGenerator(base.BaseEstimator):
             permission if the connection isn't fully setup.
         max_iterations (Optional[int], Default to 300):
             The number of steps to run when performing supervised tuning.
-        evaluation_task (Optional[str], default to "UNSPECIFIED"):
-            When performing supervised tuning, the type of task that you want to tune the model to perform. Possible values:
-            "TEXT_GENERATION", "CLASSIFICATION", "SUMMARIZATION", "QUESTION_ANSWERING", "UNSPECIFIED". Default to "UNSPECIFIED".
     """
 
     def __init__(
@@ -87,18 +83,10 @@ def __init__(
         session: Optional[bigframes.Session] = None,
         connection_name: Optional[str] = None,
         max_iterations: int = 300,
-        evaluation_task: Literal[
-            "UNSPECIFIED",
-            "TEXT_GENERATION",
-            "CLASSIFICATION",
-            "SUMMARIZATION",
-            "QUESTION_ANSWERING",
-        ] = "UNSPECIFIED",
     ):
         self.model_name = model_name
         self.session = session or bpd.get_global_session()
         self.max_iterations = max_iterations
-        self.evaluation_task = evaluation_task
         self._bq_connection_manager = self.session.bqconnectionmanager
 
         connection_name = connection_name or self.session._bq_connection
@@ -169,8 +157,6 @@ def _from_bq(
                 # Convert types
                 if bf_param in ["max_iterations"]:
                     kwargs[bf_param] = int(last_fitting[bqml_param])
-                elif bf_param in ["evaluation_task"]:
-                    kwargs[bf_param] = str(last_fitting[bqml_param])
 
         text_generator_model = cls(
             **kwargs,
@@ -187,7 +173,6 @@ def _bqml_options(self) -> dict:
         options = {
             "max_iterations": self.max_iterations,
             "data_split_method": "NO_SPLIT",
-            "evaluation_task": self.evaluation_task,
         }
         return options
 
@@ -209,7 +194,6 @@ def fit(
         """
         X, y = utils.convert_to_dataframe(X, y)
 
-        # TODO(ashleyxu): options= self._bqml_options
         options = self._bqml_options
         options["endpoint"] = self.model_name + "@001"
         options["prompt_col"] = X.columns.tolist()[0]

bigframes/ml/sql.py

@@ -177,7 +177,6 @@ def create_model(
         parts.append(f"AS {source_sql}")
         return "\n".join(parts)
 
-    # Model create and alter
     def create_llm_remote_model(
         self,
         source_df: bpd.DataFrame,

tests/system/conftest.py

@@ -537,33 +537,6 @@ def penguins_df_default_index(
     return session.read_gbq(penguins_table_id)
 
 
-@pytest.fixture(scope="session")
-def llm_fine_tune_df_default_index(
-    session: bigframes.Session,
-) -> bigframes.dataframe.DataFrame:
-    sql = """
-SELECT
-  CONCAT("Please do sentiment analysis on the following text and only output a number from 0 to 5 where 0 means sadness, 1 means joy, 2 means love, 3 means anger, 4 means fear, and 5 means surprise. Text: ", text) as prompt,
-  CAST(label AS STRING) as label
-FROM `llm_tuning.emotion_classification_train`
-"""
-    return session.read_gbq(sql)
-
-
-@pytest.fixture(scope="session")
-def llm_remote_text_pandas_df():
-    """Additional data matching the penguins dataset, with a new index"""
-    return pd.DataFrame(
-        {
-            "prompt": [
-                "Please do sentiment analysis on the following text and only output a number from 0 to 5where 0 means sadness, 1 means joy, 2 means love, 3 means anger, 4 means fear, and 5 means surprise. Text: i feel beautifully emotional knowing that these women of whom i knew just a handful were holding me and my baba on our journey",
-                "Please do sentiment analysis on the following text and only output a number from 0 to 5 where 0 means sadness, 1 means joy, 2 means love, 3 means anger, 4 means fear, and 5 means surprise. Text: i was feeling a little vain when i did this one",
-                "Please do sentiment analysis on the following text and only output a number from 0 to 5 where 0 means sadness, 1 means joy, 2 means love, 3 means anger, 4 means fear, and 5 means surprise. Text: a father of children killed in an accident",
-            ],
-        }
-    )
-
-
 @pytest.fixture(scope="session")
 def time_series_df_default_index(
     time_series_table_id: str, session: bigframes.Session

tests/system/load/test_llm.py

@@ -12,14 +12,44 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import pandas as pd
+import pytest
+
 import bigframes.ml.llm
 
 
+@pytest.fixture(scope="session")
+def llm_fine_tune_df_default_index(
+    session: bigframes.Session,
+) -> bigframes.dataframe.DataFrame:
+    sql = """
+SELECT
+  CONCAT("Please do sentiment analysis on the following text and only output a number from 0 to 5 where 0 means sadness, 1 means joy, 2 means love, 3 means anger, 4 means fear, and 5 means surprise. Text: ", text) as prompt,
+  CAST(label AS STRING) as label
+FROM `llm_tuning.emotion_classification_train`
+"""
+    return session.read_gbq(sql)
+
+
+@pytest.fixture(scope="session")
+def llm_remote_text_pandas_df():
+    """Additional data matching the penguins dataset, with a new index"""
+    return pd.DataFrame(
+        {
+            "prompt": [
+                "Please do sentiment analysis on the following text and only output a number from 0 to 5where 0 means sadness, 1 means joy, 2 means love, 3 means anger, 4 means fear, and 5 means surprise. Text: i feel beautifully emotional knowing that these women of whom i knew just a handful were holding me and my baba on our journey",
+                "Please do sentiment analysis on the following text and only output a number from 0 to 5 where 0 means sadness, 1 means joy, 2 means love, 3 means anger, 4 means fear, and 5 means surprise. Text: i was feeling a little vain when i did this one",
+                "Please do sentiment analysis on the following text and only output a number from 0 to 5 where 0 means sadness, 1 means joy, 2 means love, 3 means anger, 4 means fear, and 5 means surprise. Text: a father of children killed in an accident",
+            ],
+        }
+    )
+
+
 def test_llm_palm_configure_fit(
     llm_fine_tune_df_default_index, llm_remote_text_pandas_df
 ):
     model = bigframes.ml.llm.PaLM2TextGenerator(
-        model_name="text-bison", max_iterations=1, evaluation_task="CLASSIFICATION"
+        model_name="text-bison", max_iterations=1
     )
 
     df = llm_fine_tune_df_default_index.dropna()
@@ -35,4 +65,4 @@ def test_llm_palm_configure_fit(
     series = df["ml_generate_text_llm_result"]
     assert all(series.str.len() == 1)
 
-    # TODO(ashleyxu): After bqml rolled out version control: save, load, check parameters to ensure configuration was kept
+    # TODO(ashleyxu b/335492787): After bqml rolled out version control: save, load, check parameters to ensure configuration was kept

tests/system/small/ml/test_llm.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Google LLC
+# Copyright 2024 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.