maxframe 1.3.1__cp38-cp38-win_amd64.whl → 2.0.0b1__cp38-cp38-win_amd64.whl

maxframe/learn/contrib/xgboost/core.py

@@ -19,26 +19,19 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import numpy as np
 
+from ....errors import TileableNotExecutedError
+
 try:
     import xgboost
 except ImportError:
     xgboost = None
 
 from ....core import OutputType
-from ...core import Model, ModelData
-from ..models import ModelApplyChunk, to_remote_model
+from ..models import ModelApplyChunk, ModelWithEval, ModelWithEvalData, to_remote_model
 from .dmatrix import DMatrix
 
 
-class BoosterData(ModelData):
-    __slots__ = ("_evals_result",)
-
-    _evals_result: Dict
-
-    def __init__(self, *args, evals_result=None, **kwargs):
-        super().__init__(*args, **kwargs)
-        self._evals_result = evals_result if evals_result is not None else dict()
-
+class BoosterData(ModelWithEvalData):
     @staticmethod
     def _get_booster_score(bst, fmap=None, importance_type="weight"):
         if not fmap:
@@ -66,13 +59,6 @@ class BoosterData(ModelData):
                 fmap_data = fmap_file.read()
         return op(self, [{}], fmap=fmap_data, importance_type=importance_type)[0]
 
-    def execute(self, session=None, **kw):
-        # The evals_result should be fetched when BoosterData.execute() is called.
-        result = super().execute(session=session, **kw)
-        if self.op.has_evals_result and self.key == self.op.outputs[0].key:
-            self._evals_result.update(self.op.outputs[1].fetch(session=session))
-        return result
-
     def predict(
         self,
         data,
@@ -103,7 +89,7 @@ class BoosterData(ModelData):
         )
 
 
-class Booster(Model):
+class Booster(ModelWithEval):
     pass
 
 
@@ -137,39 +123,64 @@ else:
             return names
 
         def __repr__(self):
-            local_model = self.fetch()
-            return repr(local_model)
+            try:
+                local_model = self.fetch()
+                return repr(local_model)
+            except TileableNotExecutedError:
+                return f"<{type(self).__name__} unexecuted>"
 
         def fit(
             self,
             X,
             y,
-            sample_weights=None,
+            sample_weight=None,
+            base_margin=None,
             eval_set=None,
+            xgb_model=None,
             sample_weight_eval_set=None,
+            base_margin_eval_set=None,
             **kw,
         ):
-            """
-            Fit the regressor. Note that fit() is an eager-execution
-            API. The call will be blocked until training finished.
+            from .train import train
 
-            Parameters
-            ----------
-            X : array_like
-                Feature matrix
-            y : array_like
-                Labels
-            sample_weight : array_like
-                instance weights
-            eval_set : list, optional
-                A list of (X, y) tuple pairs to use as validation sets, for which
-                metrics will be computed.
-                Validation metrics will help us track the performance of the model.
-            sample_weight_eval_set : list, optional
-                A list of the form [L_1, L_2, ..., L_n], where each L_i is a list
-                of group weights on the i-th validation set.
-            """
-            raise NotImplementedError
+            session = kw.pop("session", None)
+            run_kwargs = kw.pop("run_kwargs", dict())
+
+            self._n_features_in = X.shape[1]
+
+            dtrain, evals = wrap_evaluation_matrices(
+                None,
+                X,
+                y,
+                sample_weight,
+                base_margin,
+                eval_set,
+                sample_weight_eval_set,
+                base_margin_eval_set,
+            )
+            params = self.get_xgb_params()
+            if not params.get("objective"):
+                params["objective"] = "reg:squarederror"
+            self.evals_result_ = dict()
+            train_kw = {}
+
+            if getattr(self, "n_classes_", None):
+                train_kw["num_class"] = self.n_classes_
+
+            result = train(
+                params,
+                dtrain,
+                num_boost_round=self.get_num_boosting_rounds(),
+                evals=evals,
+                evals_result=self.evals_result_,
+                xgb_model=xgb_model,
+                callbacks=self.callbacks,
+                session=session,
+                run_kwargs=run_kwargs,
+                **train_kw,
+            )
+            self._Booster = result
+            return self
 
         def predict(self, data, **kw):
             """

maxframe/learn/contrib/xgboost/dmatrix.py

@@ -12,14 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import List
 
 from .... import opcodes
+from ....core import EntityData
 from ....core.entity.output_types import get_output_types
 from ....core.operator.base import Operator
 from ....core.operator.core import TileableOperatorMixin
 from ....dataframe.core import DATAFRAME_TYPE
-from ....serialization.serializables import Float64Field, KeyField, ListField
-from ....serialization.serializables.field import AnyField, BoolField, Int64Field
+from ....serialization.serializables import BoolField, Float64Field, KeyField, ListField
+from ....serialization.serializables.field import AnyField, Int64Field
 from ....tensor import tensor as astensor
 from ....tensor.core import TENSOR_TYPE
 from ....typing_ import TileableType
@@ -42,24 +44,29 @@ class ToDMatrix(Operator, TileableOperatorMixin):
     qid = AnyField("qid", default=None)
     label_lower_bound = AnyField("label_lower_bound", default=None)
     label_upper_bound = AnyField("label_upper_bound", default=None)
+    # if to collocate the data, label and weight
+    collocate = BoolField("collocate", default=False)
     enable_categorical = BoolField("enable_categorical", default=None)
 
     @property
     def output_limit(self):
+        if self.collocate:
+            return 1 + sum(bool(x) for x in [self.label, self.weight, self.base_margin])
         return 1
 
-    def _set_inputs(self, inputs):
-        super()._set_inputs(inputs)
-        if self.data is not None:
-            self.data = self._inputs[0]
-        has_label = self.label is not None
+    @classmethod
+    def _set_inputs(cls, op: "ToDMatrix", inputs: List[EntityData]):
+        super()._set_inputs(op, inputs)
+        if op.data is not None:
+            op.data = op._inputs[0]
+        has_label = op.label is not None
         if has_label:
-            self.label = self._inputs[1]
-        if self.weight is not None:
+            op.label = op._inputs[1]
+        if op.weight is not None:
             i = 1 if not has_label else 2
-            self.weight = self._inputs[i]
-        if self.base_margin is not None:
-            self.base_margin = self._inputs[-1]
+            op.weight = op._inputs[i]
+        if op.base_margin is not None:
+            op.base_margin = op._inputs[-1]
 
     @staticmethod
     def _get_kw(obj):

maxframe/learn/contrib/xgboost/predict.py

@@ -12,10 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import List
 
 import numpy as np
 
 from .... import opcodes
+from ....core import EntityData
 from ....core.entity.output_types import OutputType
 from ....core.operator.base import Operator
 from ....core.operator.core import TileableOperatorMixin
@@ -45,21 +47,24 @@ class XGBPredict(Operator, TileableOperatorMixin):
     training = BoolField("training", default=False)
     iteration_range = TupleField("iteration_range", default_factory=lambda x: (0, 0))
     strict_shape = BoolField("strict_shape", default=False)
-    flag = BoolField("flag", default=False)
 
     def __init__(self, output_types=None, gpu=None, **kw):
         super().__init__(_output_types=output_types, gpu=gpu, **kw)
 
-    def _set_inputs(self, inputs):
-        super()._set_inputs(inputs)
-        self.data = self._inputs[0]
-        self.model = self._inputs[1]
+    def has_custom_code(self) -> bool:
+        return True
+
+    @classmethod
+    def _set_inputs(cls, op: "XGBPredict", inputs: List[EntityData]):
+        super()._set_inputs(op, inputs)
+        op.data = op._inputs[0]
+        op.model = op._inputs[1]
 
     def __call__(self):
         num_class = getattr(self.model.op, "num_class", None)
         if num_class is not None:
             num_class = int(num_class)
-        if num_class is not None:
+        if num_class is not None and num_class > 2:
             shape = (self.data.shape[0], num_class)
         else:
             shape = (self.data.shape[0],)
@@ -84,7 +89,7 @@ def predict(
     training=False,
     iteration_range=None,
     strict_shape=False,
-    flag=False,
+    **kwargs,
 ):
     """
     Using MaxFrame XGBoost model to predict data.
@@ -100,7 +105,9 @@ def predict(
     import xgboost
 
     data = check_data(data)
-    if isinstance(model, xgboost.Booster):
+    if not isinstance(model, (Booster, BoosterData, xgboost.Booster)):
+        raise TypeError(f"model has to be a xgboost.Booster, got {type(model)} instead")
+    elif isinstance(model, xgboost.Booster):
         model = to_remote_model(model, model_cls=Booster)
 
     output_types = [OutputType.tensor]
@@ -120,5 +127,5 @@ def predict(
         strict_shape=strict_shape,
         gpu=data.op.gpu,
         output_types=output_types,
-        flag=flag,
+        **kwargs,
     )()

maxframe/learn/contrib/xgboost/regressor.py

@@ -22,9 +22,7 @@ if not xgboost:
 else:
     from xgboost.sklearn import XGBRegressorBase
 
-    from .core import wrap_evaluation_matrices
     from .predict import predict
-    from .train import train
 
     class XGBRegressor(XGBScikitLearnBase, XGBRegressorBase):
         """
@@ -46,39 +44,42 @@ else:
             sample_weight=None,
             base_margin=None,
             eval_set=None,
+            xgb_model=None,
             sample_weight_eval_set=None,
             base_margin_eval_set=None,
             **kw,
         ):
-            session = kw.pop("session", None)
-            run_kwargs = kw.pop("run_kwargs", dict())
+            """
+            Fit the regressor. Note that fit() is an eager-execution
+            API. The call will be blocked until training finished.
 
-            self._n_features_in = X.shape[1]
-
-            dtrain, evals = wrap_evaluation_matrices(
-                None,
+            Parameters
+            ----------
+            X : array_like
+                Feature matrix
+            y : array_like
+                Labels
+            sample_weight : array_like
+                instance weights
+            eval_set : list, optional
+                A list of (X, y) tuple pairs to use as validation sets, for which
+                metrics will be computed.
+                Validation metrics will help us track the performance of the model.
+            sample_weight_eval_set : list, optional
+                A list of the form [L_1, L_2, ..., L_n], where each L_i is a list
+                of group weights on the i-th validation set.
+            """
+            super().fit(
                 X,
                 y,
-                sample_weight,
-                base_margin,
-                eval_set,
-                sample_weight_eval_set,
-                base_margin_eval_set,
-            )
-            params = self.get_xgb_params()
-            if not params.get("objective"):
-                params["objective"] = "reg:squarederror"
-            self.evals_result_ = dict()
-            result = train(
-                params,
-                dtrain,
-                num_boost_round=self.get_num_boosting_rounds(),
-                evals=evals,
-                evals_result=self.evals_result_,
-                session=session,
-                run_kwargs=run_kwargs,
+                sample_weight=sample_weight,
+                base_margin=base_margin,
+                eval_set=eval_set,
+                xgb_model=xgb_model,
+                sample_weight_eval_set=sample_weight_eval_set,
+                base_margin_eval_set=base_margin_eval_set,
+                **kw,
             )
-            self._Booster = result
             return self
 
         def predict(self, data, **kw):

maxframe/learn/contrib/xgboost/tests/test_callback.py

@@ -0,0 +1,41 @@
+# Copyright 1999-2025 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+from ..callback import EarlyStopping, LearningRateScheduler, XGBTrainingCallback
+
+try:
+    from xgboost.callback import EarlyStopping as _EarlyStopping
+except ImportError:
+    pytestmark = pytest.mark.skip("Need xgboost to run the test")
+
+
+def test_xgb_training_callbacks():
+    early_stopping_cb = EarlyStopping(rounds=5)
+    assert not early_stopping_cb.has_custom_code()
+
+    local_early_stopping_cb = early_stopping_cb.to_local()
+    assert isinstance(local_early_stopping_cb, _EarlyStopping)
+    assert local_early_stopping_cb.rounds == 5
+
+    rmt_early_stopping_cb = XGBTrainingCallback.from_local(local_early_stopping_cb)
+    assert isinstance(rmt_early_stopping_cb, EarlyStopping)
+    assert rmt_early_stopping_cb.rounds == 5
+
+    lr_sch_cb = LearningRateScheduler([0.5, 0.4, 0.1])
+    assert not lr_sch_cb.has_custom_code()
+
+    lr_sch_cb = LearningRateScheduler(lambda x: x * 0.1)
+    assert lr_sch_cb.has_custom_code()

maxframe/learn/contrib/xgboost/train.py

@@ -14,11 +14,11 @@
 
 import logging
 from collections import OrderedDict
+from typing import List
 
 from .... import opcodes
-from ....core import OutputType
-from ....core.operator.base import Operator
-from ....core.operator.core import TileableOperatorMixin
+from ....core import EntityData, OutputType
+from ....core.operator import ObjectOperator, ObjectOperatorMixin
 from ....serialization.serializables import (
     AnyField,
     BoolField,
@@ -29,7 +29,9 @@ from ....serialization.serializables import (
     KeyField,
     ListField,
 )
-from .core import Booster
+from ..models import to_remote_model
+from ..utils import TrainingCallback
+from .core import Booster, BoosterData, XGBScikitLearnBase
 from .dmatrix import ToDMatrix, to_dmatrix
 
 logger = logging.getLogger(__name__)
@@ -41,7 +43,7 @@ def _on_serialize_evals(evals_val):
     return [list(x) for x in evals_val]
 
 
-class XGBTrain(Operator, TileableOperatorMixin):
+class XGBTrain(ObjectOperator, ObjectOperatorMixin):
     _op_type_ = opcodes.XGBOOST_TRAIN
 
     params = DictField("params", key_type=FieldTypes.string, default=None)
@@ -52,9 +54,12 @@ class XGBTrain(Operator, TileableOperatorMixin):
     maximize = BoolField("maximize", default=None)
     early_stopping_rounds = Int64Field("early_stopping_rounds", default=None)
     verbose_eval = AnyField("verbose_eval", default=None)
-    xgb_model = AnyField("xgb_model", default=None)
+    xgb_model = KeyField("xgb_model", default=None)
     callbacks = ListField(
-        "callbacks", field_type=FunctionField.field_type, default=None
+        "callbacks",
+        field_type=FunctionField.field_type,
+        default=None,
+        on_serialize=TrainingCallback.from_local,
     )
     custom_metric = FunctionField("custom_metric", default=None)
     num_boost_round = Int64Field("num_boost_round", default=10)
@@ -67,21 +72,35 @@ class XGBTrain(Operator, TileableOperatorMixin):
         if self.has_evals_result:
             self.output_types.append(OutputType.object)
 
-    def _set_inputs(self, inputs):
-        super()._set_inputs(inputs)
-        self.dtrain = self._inputs[0]
-        rest = self._inputs[1:]
-        if self.has_evals_result:
-            evals_dict = OrderedDict(self.evals)
+    def has_custom_code(self) -> bool:
+        if not self.callbacks:
+            return False
+        return any(
+            not isinstance(cb, TrainingCallback) or cb.has_custom_code()
+            for cb in self.callbacks
+        )
+
+    @classmethod
+    def _set_inputs(cls, op: "XGBTrain", inputs: List[EntityData]):
+        super()._set_inputs(op, inputs)
+        input_it = iter(op._inputs)
+        op.dtrain = next(input_it)
+        if op.evals:
+            evals_dict = OrderedDict(op.evals)
             new_evals_dict = OrderedDict()
-            for new_key, val in zip(rest, evals_dict.values()):
+            for val in evals_dict.values():
+                new_key = next(input_it)
                 new_evals_dict[new_key] = val
-            self.evals = list(new_evals_dict.items())
+            op.evals = list(new_evals_dict.items())
+        if op.xgb_model:
+            op.xgb_model = next(input_it)
 
     def __call__(self, evals_result):
         inputs = [self.dtrain]
         if self.has_evals_result:
             inputs.extend(e[0] for e in self.evals)
+        if self.xgb_model is not None:
+            inputs.append(self.xgb_model)
         kws = [{"object_class": Booster}, {}]
         return self.new_tileables(inputs, kws=kws, evals_result=evals_result)[0]
 
@@ -94,7 +113,28 @@ class XGBTrain(Operator, TileableOperatorMixin):
         return self.evals
 
 
-def train(params, dtrain, evals=None, evals_result=None, num_class=None, **kwargs):
+def _get_xgb_booster(xgb_model):
+    import xgboost
+
+    if isinstance(xgb_model, (XGBScikitLearnBase, xgboost.XGBModel)):
+        xgb_model = xgb_model.get_booster()
+
+    if isinstance(xgb_model, (Booster, BoosterData)):
+        return xgb_model
+    elif isinstance(xgb_model, xgboost.Booster):
+        return to_remote_model(xgb_model, model_cls=Booster)
+    raise ValueError(f"Cannot use {type(xgb_model)} as xgb_model")
+
+
+def train(
+    params,
+    dtrain,
+    evals=None,
+    evals_result=None,
+    xgb_model=None,
+    num_class=None,
+    **kwargs,
+):
     """
     Train XGBoost model in MaxFrame manner.
 
@@ -120,11 +160,14 @@ def train(params, dtrain, evals=None, evals_result=None, num_class=None, **kwarg
                 processed_evals.append((eval_dmatrix, name))
             else:
                 processed_evals.append((to_dmatrix(eval_dmatrix), name))
+    if xgb_model is not None:
+        xgb_model = _get_xgb_booster(xgb_model)
     data = XGBTrain(
         params=params,
         dtrain=dtrain,
         evals=processed_evals,
         evals_result=evals_result,
+        xgb_model=xgb_model,
         num_class=num_class,
         **kwargs,
     )(evals_result)