feat: allow fit to take additional eval data in linear and ensemble models (#1096)

* feat: allow `fit` to take additional eval data

The additional eval data would be used to measure the fitted model and
attach the measurement to the underlying BQML model, which can be used
as benchmark for the model consumers in BigQuery Studio and otherwise.

* subclass from TrainablePredictor

* add support for fit-time evaluation in ensemble models

* fetch logistic regression eval numbers from multiClassClassificationMetrics

* use the generic type template

* update vendored docstrings for fit taking X_eval, y_eval

* update key to fetch model eval metrics

* enfore binary classification in the logistic regression test

Author: shobsi

bigframes/ml/base.py

@@ -164,6 +164,40 @@ def fit(
         return self._fit(X, y)
 
 
+class TrainableWithEvaluationPredictor(TrainablePredictor):
+    """A BigQuery DataFrames ML Model base class that can be used to fit and predict outputs.
+
+    Additional evaluation data can be provided to measure the model in the fit phase."""
+
+    @abc.abstractmethod
+    def _fit(self, X, y, transforms=None, X_eval=None, y_eval=None):
+        pass
+
+    @abc.abstractmethod
+    def score(self, X, y):
+        pass
+
+
+class SupervisedTrainableWithEvaluationPredictor(TrainableWithEvaluationPredictor):
+    """A BigQuery DataFrames ML Supervised Model base class that can be used to fit and predict outputs.
+
+    Need to provide both X and y in supervised tasks.
+
+    Additional X_eval and y_eval can be provided to measure the model in the fit phase.
+    """
+
+    _T = TypeVar("_T", bound="SupervisedTrainableWithEvaluationPredictor")
+
+    def fit(
+        self: _T,
+        X: utils.ArrayType,
+        y: utils.ArrayType,
+        X_eval: Optional[utils.ArrayType] = None,
+        y_eval: Optional[utils.ArrayType] = None,
+    ) -> _T:
+        return self._fit(X, y, X_eval=X_eval, y_eval=y_eval)
+
+
 class UnsupervisedTrainablePredictor(TrainablePredictor):
     """A BigQuery DataFrames ML Unsupervised Model base class that can be used to fit and predict outputs.
 

bigframes/ml/ensemble.py

@@ -52,7 +52,7 @@
 
 @log_adapter.class_logger
 class XGBRegressor(
-    base.SupervisedTrainablePredictor,
+    base.SupervisedTrainableWithEvaluationPredictor,
     bigframes_vendored.xgboost.sklearn.XGBRegressor,
 ):
     __doc__ = bigframes_vendored.xgboost.sklearn.XGBRegressor.__doc__
@@ -145,14 +145,24 @@ def _fit(
         X: utils.ArrayType,
         y: utils.ArrayType,
         transforms: Optional[List[str]] = None,
+        X_eval: Optional[utils.ArrayType] = None,
+        y_eval: Optional[utils.ArrayType] = None,
     ) -> XGBRegressor:
         X, y = utils.convert_to_dataframe(X, y)
 
+        bqml_options = self._bqml_options
+
+        if X_eval is not None and y_eval is not None:
+            X_eval, y_eval = utils.convert_to_dataframe(X_eval, y_eval)
+            X, y, bqml_options = utils.combine_training_and_evaluation_data(
+                X, y, X_eval, y_eval, bqml_options
+            )
+
         self._bqml_model = self._bqml_model_factory.create_model(
             X,
             y,
             transforms=transforms,
-            options=self._bqml_options,
+            options=bqml_options,
         )
         return self
 
@@ -200,7 +210,7 @@ def to_gbq(self, model_name: str, replace: bool = False) -> XGBRegressor:
 
 @log_adapter.class_logger
 class XGBClassifier(
-    base.SupervisedTrainablePredictor,
+    base.SupervisedTrainableWithEvaluationPredictor,
     bigframes_vendored.xgboost.sklearn.XGBClassifier,
 ):
 
@@ -294,14 +304,24 @@ def _fit(
         X: utils.ArrayType,
         y: utils.ArrayType,
         transforms: Optional[List[str]] = None,
+        X_eval: Optional[utils.ArrayType] = None,
+        y_eval: Optional[utils.ArrayType] = None,
     ) -> XGBClassifier:
         X, y = utils.convert_to_dataframe(X, y)
 
+        bqml_options = self._bqml_options
+
+        if X_eval is not None and y_eval is not None:
+            X_eval, y_eval = utils.convert_to_dataframe(X_eval, y_eval)
+            X, y, bqml_options = utils.combine_training_and_evaluation_data(
+                X, y, X_eval, y_eval, bqml_options
+            )
+
         self._bqml_model = self._bqml_model_factory.create_model(
             X,
             y,
             transforms=transforms,
-            options=self._bqml_options,
+            options=bqml_options,
         )
         return self
 
@@ -347,7 +367,7 @@ def to_gbq(self, model_name: str, replace: bool = False) -> XGBClassifier:
 
 @log_adapter.class_logger
 class RandomForestRegressor(
-    base.SupervisedTrainablePredictor,
+    base.SupervisedTrainableWithEvaluationPredictor,
     bigframes_vendored.sklearn.ensemble._forest.RandomForestRegressor,
 ):
 
@@ -430,14 +450,24 @@ def _fit(
         X: utils.ArrayType,
         y: utils.ArrayType,
         transforms: Optional[List[str]] = None,
+        X_eval: Optional[utils.ArrayType] = None,
+        y_eval: Optional[utils.ArrayType] = None,
     ) -> RandomForestRegressor:
         X, y = utils.convert_to_dataframe(X, y)
 
+        bqml_options = self._bqml_options
+
+        if X_eval is not None and y_eval is not None:
+            X_eval, y_eval = utils.convert_to_dataframe(X_eval, y_eval)
+            X, y, bqml_options = utils.combine_training_and_evaluation_data(
+                X, y, X_eval, y_eval, bqml_options
+            )
+
         self._bqml_model = self._bqml_model_factory.create_model(
             X,
             y,
             transforms=transforms,
-            options=self._bqml_options,
+            options=bqml_options,
         )
         return self
 
@@ -503,7 +533,7 @@ def to_gbq(self, model_name: str, replace: bool = False) -> RandomForestRegresso
 
 @log_adapter.class_logger
 class RandomForestClassifier(
-    base.SupervisedTrainablePredictor,
+    base.SupervisedTrainableWithEvaluationPredictor,
     bigframes_vendored.sklearn.ensemble._forest.RandomForestClassifier,
 ):
 
@@ -586,14 +616,24 @@ def _fit(
         X: utils.ArrayType,
         y: utils.ArrayType,
         transforms: Optional[List[str]] = None,
+        X_eval: Optional[utils.ArrayType] = None,
+        y_eval: Optional[utils.ArrayType] = None,
     ) -> RandomForestClassifier:
         X, y = utils.convert_to_dataframe(X, y)
 
+        bqml_options = self._bqml_options
+
+        if X_eval is not None and y_eval is not None:
+            X_eval, y_eval = utils.convert_to_dataframe(X_eval, y_eval)
+            X, y, bqml_options = utils.combine_training_and_evaluation_data(
+                X, y, X_eval, y_eval, bqml_options
+            )
+
         self._bqml_model = self._bqml_model_factory.create_model(
             X,
             y,
             transforms=transforms,
-            options=self._bqml_options,
+            options=bqml_options,
         )
         return self
 

bigframes/ml/linear_model.py

@@ -47,7 +47,7 @@
 
 @log_adapter.class_logger
 class LinearRegression(
-    base.SupervisedTrainablePredictor,
+    base.SupervisedTrainableWithEvaluationPredictor,
     bigframes_vendored.sklearn.linear_model._base.LinearRegression,
 ):
     __doc__ = bigframes_vendored.sklearn.linear_model._base.LinearRegression.__doc__
@@ -131,14 +131,24 @@ def _fit(
         X: utils.ArrayType,
         y: utils.ArrayType,
         transforms: Optional[List[str]] = None,
+        X_eval: Optional[utils.ArrayType] = None,
+        y_eval: Optional[utils.ArrayType] = None,
     ) -> LinearRegression:
         X, y = utils.convert_to_dataframe(X, y)
 
+        bqml_options = self._bqml_options
+
+        if X_eval is not None and y_eval is not None:
+            X_eval, y_eval = utils.convert_to_dataframe(X_eval, y_eval)
+            X, y, bqml_options = utils.combine_training_and_evaluation_data(
+                X, y, X_eval, y_eval, bqml_options
+            )
+
         self._bqml_model = self._bqml_model_factory.create_model(
             X,
             y,
             transforms=transforms,
-            options=self._bqml_options,
+            options=bqml_options,
         )
         return self
 
@@ -183,7 +193,7 @@ def to_gbq(self, model_name: str, replace: bool = False) -> LinearRegression:
 
 @log_adapter.class_logger
 class LogisticRegression(
-    base.SupervisedTrainablePredictor,
+    base.SupervisedTrainableWithEvaluationPredictor,
     bigframes_vendored.sklearn.linear_model._logistic.LogisticRegression,
 ):
     __doc__ = (
@@ -283,15 +293,24 @@ def _fit(
         X: utils.ArrayType,
         y: utils.ArrayType,
         transforms: Optional[List[str]] = None,
+        X_eval: Optional[utils.ArrayType] = None,
+        y_eval: Optional[utils.ArrayType] = None,
     ) -> LogisticRegression:
-        """Fit model with transforms."""
         X, y = utils.convert_to_dataframe(X, y)
 
+        bqml_options = self._bqml_options
+
+        if X_eval is not None and y_eval is not None:
+            X_eval, y_eval = utils.convert_to_dataframe(X_eval, y_eval)
+            X, y, bqml_options = utils.combine_training_and_evaluation_data(
+                X, y, X_eval, y_eval, bqml_options
+            )
+
         self._bqml_model = self._bqml_model_factory.create_model(
             X,
             y,
             transforms=transforms,
-            options=self._bqml_options,
+            options=bqml_options,
         )
         return self
 

bigframes/ml/utils.py

@@ -13,13 +13,13 @@
 # limitations under the License.
 
 import typing
-from typing import Any, Generator, Literal, Mapping, Optional, Union
+from typing import Any, Generator, Literal, Mapping, Optional, Tuple, Union
 
 import bigframes_vendored.constants as constants
 from google.cloud import bigquery
 import pandas as pd
 
-from bigframes.core import blocks
+from bigframes.core import blocks, guid
 import bigframes.pandas as bpd
 from bigframes.session import Session
 
@@ -155,3 +155,37 @@ def retrieve_params_from_bq_model(
             kwargs[bf_param] = bf_param_type(last_fitting[bqml_param])
 
     return kwargs
+
+
+def combine_training_and_evaluation_data(
+    X_train: bpd.DataFrame,
+    y_train: bpd.DataFrame,
+    X_eval: bpd.DataFrame,
+    y_eval: bpd.DataFrame,
+    bqml_options: dict,
+) -> Tuple[bpd.DataFrame, bpd.DataFrame, dict]:
+    """
+    Combine training data and labels with evlauation data and labels, and keep
+    them differentiated through a split column in the combined data and labels.
+    """
+
+    assert X_train.columns.equals(X_eval.columns)
+    assert y_train.columns.equals(y_eval.columns)
+
+    # create a custom split column for BQML and supply the evaluation
+    # data along with the training data in a combined single table
+    # https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-create-dnn-models#data_split_col.
+    split_col = guid.generate_guid()
+    assert split_col not in X_train.columns
+
+    X_train[split_col] = False
+    X_eval[split_col] = True
+    X = bpd.concat([X_train, X_eval])
+    y = bpd.concat([y_train, y_eval])
+
+    # create options copy to not mutate the incoming one
+    bqml_options = bqml_options.copy()
+    bqml_options["data_split_method"] = "CUSTOM"
+    bqml_options["data_split_col"] = split_col
+
+    return X, y, bqml_options