feat: add support for LinearRegression.predict_explain and LogisticRegression.predict_explain parameter, top_k_features (#1228)

* feat: add LogisticRegression.predict_explain() to generate ML.EXPLAIN_PREDICT columns

* update tests

* chore: add support for predict_explain paramater, top_k_features

* update test

* update logistic reg method with the new param

* add and test new param's validation

* Update bigframes/ml/linear_model.py

---------

Co-authored-by: Tim Sweña (Swast) <swast@google.com>

Author: arwas11

bigframes/ml/core.py

@@ -123,10 +123,15 @@ def predict(self, input_data: bpd.DataFrame) -> bpd.DataFrame:
             self._model_manipulation_sql_generator.ml_predict,
         )
 
-    def explain_predict(self, input_data: bpd.DataFrame) -> bpd.DataFrame:
+    def explain_predict(
+        self, input_data: bpd.DataFrame, options: Mapping[str, int | float]
+    ) -> bpd.DataFrame:
         return self._apply_ml_tvf(
             input_data,
-            self._model_manipulation_sql_generator.ml_explain_predict,
+            lambda source_sql: self._model_manipulation_sql_generator.ml_explain_predict(
+                source_sql=source_sql,
+                struct_options=options,
+            ),
         )
 
     def transform(self, input_data: bpd.DataFrame) -> bpd.DataFrame:

bigframes/ml/linear_model.py

@@ -155,14 +155,15 @@ def _fit(
     def predict(self, X: utils.ArrayType) -> bpd.DataFrame:
         if not self._bqml_model:
             raise RuntimeError("A model must be fitted before predict")
-
-        (X,) = utils.batch_convert_to_dataframe(X)
+        (X,) = utils.batch_convert_to_dataframe(X, session=self._bqml_model.session)
 
         return self._bqml_model.predict(X)
 
     def predict_explain(
         self,
         X: utils.ArrayType,
+        *,
+        top_k_features: int = 5,
     ) -> bpd.DataFrame:
         """
         Explain predictions for a linear regression model.
@@ -175,18 +176,32 @@ def predict_explain(
             X (bigframes.dataframe.DataFrame or bigframes.series.Series or
             pandas.core.frame.DataFrame or pandas.core.series.Series):
                 Series or a DataFrame to explain its predictions.
+            top_k_features (int, default 5):
+                an INT64 value that specifies how many top feature attribution
+                pairs are generated for each row of input data. The features are
+                ranked by the absolute values of their attributions.
+
+                By default, top_k_features is set to 5. If its value is greater
+                than the number of features in the training data, the
+                attributions of all features are returned.
 
         Returns:
             bigframes.pandas.DataFrame:
                 The predicted DataFrames with explanation columns.
         """
-        # TODO(b/377366612): Add support for `top_k_features` parameter
+        if top_k_features < 1:
+            raise ValueError(
+                f"top_k_features must be at least 1, but is {top_k_features}."
+            )
+
         if not self._bqml_model:
             raise RuntimeError("A model must be fitted before predict")
 
         (X,) = utils.batch_convert_to_dataframe(X, session=self._bqml_model.session)
 
-        return self._bqml_model.explain_predict(X)
+        return self._bqml_model.explain_predict(
+            X, options={"top_k_features": top_k_features}
+        )
 
     def score(
         self,
@@ -356,6 +371,8 @@ def predict(
     def predict_explain(
         self,
         X: utils.ArrayType,
+        *,
+        top_k_features: int = 5,
     ) -> bpd.DataFrame:
         """
         Explain predictions for a logistic regression model.
@@ -368,18 +385,32 @@ def predict_explain(
             X (bigframes.dataframe.DataFrame or bigframes.series.Series or
             pandas.core.frame.DataFrame or pandas.core.series.Series):
                 Series or a DataFrame to explain its predictions.
+            top_k_features (int, default 5):
+                an INT64 value that specifies how many top feature attribution
+                pairs are generated for each row of input data. The features are
+                ranked by the absolute values of their attributions.
+
+                By default, top_k_features is set to 5. If its value is greater
+                than the number of features in the training data, the
+                attributions of all features are returned.
 
         Returns:
             bigframes.pandas.DataFrame:
                 The predicted DataFrames with explanation columns.
         """
-        # TODO(b/377366612): Add support for `top_k_features` parameter
+        if top_k_features < 1:
+            raise ValueError(
+                f"top_k_features must be at least 1, but is {top_k_features}."
+            )
+
         if not self._bqml_model:
             raise RuntimeError("A model must be fitted before predict")
 
         (X,) = utils.batch_convert_to_dataframe(X, session=self._bqml_model.session)
 
-        return self._bqml_model.explain_predict(X)
+        return self._bqml_model.explain_predict(
+            X, options={"top_k_features": top_k_features}
+        )
 
     def score(
         self,

bigframes/ml/sql.py

@@ -304,10 +304,13 @@ def ml_predict(self, source_sql: str) -> str:
         return f"""SELECT * FROM ML.PREDICT(MODEL {self._model_ref_sql()},
   ({source_sql}))"""
 
-    def ml_explain_predict(self, source_sql: str) -> str:
+    def ml_explain_predict(
+        self, source_sql: str, struct_options: Mapping[str, Union[int, float]]
+    ) -> str:
         """Encode ML.EXPLAIN_PREDICT for BQML"""
+        struct_options_sql = self.struct_options(**struct_options)
         return f"""SELECT * FROM ML.EXPLAIN_PREDICT(MODEL {self._model_ref_sql()},
-  ({source_sql}))"""
+  ({source_sql}), {struct_options_sql})"""
 
     def ml_forecast(self, struct_options: Mapping[str, Union[int, float]]) -> str:
         """Encode ML.FORECAST for BQML"""

tests/system/small/ml/test_core.py

@@ -263,8 +263,9 @@ def test_model_predict(penguins_bqml_linear_model: core.BqmlModel, new_penguins_
 def test_model_predict_explain(
     penguins_bqml_linear_model: core.BqmlModel, new_penguins_df
 ):
+    options = {"top_k_features": 3}
     predictions = penguins_bqml_linear_model.explain_predict(
-        new_penguins_df
+        new_penguins_df, options
     ).to_pandas()
     expected = pd.DataFrame(
         {
@@ -317,14 +318,15 @@ def test_model_predict_explain_with_unnamed_index(
     # need to persist through the call to ML.PREDICT
     new_penguins_df = new_penguins_df.reset_index()
 
+    options = {"top_k_features": 3}
     # remove the middle tag number to ensure we're really keeping the unnamed index
     new_penguins_df = typing.cast(
         bigframes.dataframe.DataFrame,
         new_penguins_df[new_penguins_df.tag_number != 1672],
     )
 
     predictions = penguins_bqml_linear_model.explain_predict(
-        new_penguins_df
+        new_penguins_df, options
     ).to_pandas()
 
     expected = pd.DataFrame(

tests/system/small/ml/test_linear_model.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import re
+
 import google.api_core.exceptions
 import pandas
 import pytest
@@ -132,6 +134,20 @@ def test_linear_reg_model_predict_explain(penguins_linear_model, new_penguins_df
     )
 
 
+def test_linear_model_predict_explain_top_k_features(
+    penguins_logistic_model: linear_model.LinearRegression, new_penguins_df
+):
+    top_k_features = 0
+
+    with pytest.raises(
+        ValueError,
+        match=re.escape(f"top_k_features must be at least 1, but is {top_k_features}."),
+    ):
+        penguins_logistic_model.predict_explain(
+            new_penguins_df, top_k_features=top_k_features
+        ).to_pandas()
+
+
 def test_linear_reg_model_predict_params(
     penguins_linear_model: linear_model.LinearRegression, new_penguins_df
 ):
@@ -307,6 +323,20 @@ def test_logistic_model_predict(penguins_logistic_model, new_penguins_df):
     )
 
 
+def test_logistic_model_predict_explain_top_k_features(
+    penguins_logistic_model: linear_model.LogisticRegression, new_penguins_df
+):
+    top_k_features = 0
+
+    with pytest.raises(
+        ValueError,
+        match=re.escape(f"top_k_features must be at least 1, but is {top_k_features}."),
+    ):
+        penguins_logistic_model.predict_explain(
+            new_penguins_df, top_k_features=top_k_features
+        ).to_pandas()
+
+
 def test_logistic_model_predict_params(
     penguins_logistic_model: linear_model.LogisticRegression, new_penguins_df
 ):

tests/unit/ml/test_sql.py

@@ -342,18 +342,6 @@ def test_ml_predict_correct(
     )
 
 
-def test_ml_explain_predict_correct(
-    model_manipulation_sql_generator: ml_sql.ModelManipulationSqlGenerator,
-    mock_df: bpd.DataFrame,
-):
-    sql = model_manipulation_sql_generator.ml_explain_predict(source_sql=mock_df.sql)
-    assert (
-        sql
-        == """SELECT * FROM ML.EXPLAIN_PREDICT(MODEL `my_project_id`.`my_dataset_id`.`my_model_id`,
-  (input_X_y_sql))"""
-    )
-
-
 def test_ml_llm_evaluate_correct(
     model_manipulation_sql_generator: ml_sql.ModelManipulationSqlGenerator,
     mock_df: bpd.DataFrame,
@@ -462,6 +450,23 @@ def test_ml_generate_embedding_correct(
     )
 
 
+def test_ml_explain_predict_correct(
+    model_manipulation_sql_generator: ml_sql.ModelManipulationSqlGenerator,
+    mock_df: bpd.DataFrame,
+):
+    sql = model_manipulation_sql_generator.ml_explain_predict(
+        source_sql=mock_df.sql,
+        struct_options={"option_key1": 1, "option_key2": 2.25},
+    )
+    assert (
+        sql
+        == """SELECT * FROM ML.EXPLAIN_PREDICT(MODEL `my_project_id`.`my_dataset_id`.`my_model_id`,
+  (input_X_y_sql), STRUCT(
+  1 AS `option_key1`,
+  2.25 AS `option_key2`))"""
+    )
+
+
 def test_ml_detect_anomalies_correct_sql(
     model_manipulation_sql_generator: ml_sql.ModelManipulationSqlGenerator,
     mock_df: bpd.DataFrame,