maxframe 2.4.0rc1__cp312-cp312-win32.whl

maxframe/learn/contrib/xgboost/dmatrix.py

@@ -0,0 +1,157 @@
+# 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.
+
+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 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
+from ...utils import convert_to_tensor_or_dataframe
+
+
+class ToDMatrix(Operator, TileableOperatorMixin):
+    _op_type_ = opcodes.TO_DMATRIX
+
+    data = KeyField("data", default=None)
+    label = KeyField("label", default=None)
+    missing = Float64Field("missing", default=None)
+    weight = KeyField("weight", default=None)
+    base_margin = KeyField("base_margin", default=None)
+    feature_names = ListField("feature_names", default=None)
+    feature_types = ListField("feature_types", default=None)
+    feature_weights = AnyField("feature_weights", default=None)
+    nthread = Int64Field("nthread", default=None)
+    group = AnyField("group", default=None)
+    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
+
+    @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:
+            op.label = op._inputs[1]
+        if op.weight is not None:
+            i = 1 if not has_label else 2
+            op.weight = op._inputs[i]
+        if op.base_margin is not None:
+            op.base_margin = op._inputs[-1]
+
+    @staticmethod
+    def _get_kw(obj):
+        if isinstance(obj, TENSOR_TYPE):
+            return {"shape": obj.shape, "dtype": obj.dtype, "order": obj.order}
+        else:
+            return {
+                "shape": obj.shape,
+                "dtypes": obj.dtypes,
+                "index_value": obj.index_value,
+                "columns_value": obj.columns_value,
+            }
+
+    def __call__(self):
+        inputs = [self.data]
+        kw = self._get_kw(self.data)
+        if self.label is not None:
+            inputs.append(self.label)
+        if self.weight is not None:
+            inputs.append(self.weight)
+        if self.base_margin is not None:
+            inputs.append(self.base_margin)
+
+        return self.new_tileable(inputs, **kw)
+
+
+def check_data(data):
+    data = convert_to_tensor_or_dataframe(data)
+    if data.ndim != 2:
+        raise ValueError(f"Expecting 2-d data, got: {data.ndim}-d")
+
+    return data
+
+
+def check_array_like(y: TileableType, name: str) -> TileableType:
+    if y is None:
+        return
+    y = convert_to_tensor_or_dataframe(y)
+    if isinstance(y, DATAFRAME_TYPE):
+        y = y.iloc[:, 0]
+    return astensor(y)
+
+
+def to_dmatrix(
+    data,
+    label=None,
+    missing=None,
+    weight=None,
+    base_margin=None,
+    feature_names=None,
+    feature_types=None,
+    feature_weights=None,
+    nthread=None,
+    group=None,
+    qid=None,
+    label_lower_bound=None,
+    label_upper_bound=None,
+    enable_categorical=None,
+):
+    data = check_data(data)
+    label = check_array_like(label, "label")
+    weight = check_array_like(weight, "weight")
+    base_margin = check_array_like(base_margin, "base_margin")
+
+    # If not multiple outputs, try to collect the chunks on same worker into one
+    # to feed the data into XGBoost for training.
+    op = ToDMatrix(
+        data=data,
+        label=label,
+        missing=missing,
+        weight=weight,
+        base_margin=base_margin,
+        feature_names=feature_names,
+        feature_types=feature_types,
+        feature_weights=feature_weights,
+        nthread=nthread,
+        group=group,
+        qid=qid,
+        label_lower_bound=label_lower_bound,
+        label_upper_bound=label_upper_bound,
+        gpu=data.op.gpu,
+        _output_types=get_output_types(data),
+        enable_categorical=enable_categorical,
+    )
+    return op()
+
+
+DMatrix = to_dmatrix

maxframe/learn/contrib/xgboost/predict.py

@@ -0,0 +1,133 @@
+# 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.
+
+from typing import List
+
+import numpy as np
+import pandas as pd
+
+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
+from ....serialization.serializables import (
+    BoolField,
+    KeyField,
+    ReferenceField,
+    TupleField,
+)
+from ....tensor.core import TensorOrder
+from ..models import to_remote_model
+from .core import Booster, BoosterData
+from .dmatrix import check_data
+
+
+class XGBPredict(Operator, TileableOperatorMixin):
+    _op_type_ = opcodes.XGBOOST_PREDICT
+    output_dtype = np.dtype(np.float32)
+
+    data = KeyField("data", default=None)
+    model = ReferenceField("model", reference_type=BoosterData, default=None)
+    pred_leaf = BoolField("pred_leaf", default=False)
+    pred_contribs = BoolField("pred_contribs", default=False)
+    approx_contribs = BoolField("approx_contribs", default=False)
+    pred_interactions = BoolField("pred_interactions", default=False)
+    validate_features = BoolField("validate_features", default=True)
+    training = BoolField("training", default=False)
+    iteration_range = TupleField("iteration_range", default_factory=lambda x: (0, 0))
+    strict_shape = BoolField("strict_shape", default=False)
+
+    def __init__(self, output_types=None, gpu=None, **kw):
+        super().__init__(_output_types=output_types, gpu=gpu, **kw)
+
+    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)
+        output_ndim = getattr(self.model.op, "output_ndim", None)
+        if num_class is not None and not pd.isna(num_class):
+            num_class = int(num_class)
+        if num_class is not None and (num_class > 2 or output_ndim == 2):
+            shape = (self.data.shape[0], num_class)
+        else:
+            shape = (self.data.shape[0],)
+        inputs = [self.data, self.model]
+        return self.new_tileable(
+            inputs,
+            shape=shape,
+            dtype=self.output_dtype,
+            order=TensorOrder.C_ORDER,
+        )
+
+
+def predict(
+    model,
+    data,
+    output_margin=False,
+    pred_leaf=False,
+    pred_contribs=False,
+    approx_contribs=False,
+    pred_interactions=False,
+    validate_features=True,
+    training=False,
+    iteration_range=None,
+    strict_shape=False,
+    **kwargs,
+):
+    """
+    Using MaxFrame XGBoost model to predict data.
+
+    Parameters
+    ----------
+    Parameters are the same as `xgboost.train`. The predict() is lazy-execution mode.
+
+    Returns
+    -------
+    results: Booster
+    """
+    import xgboost
+
+    data = check_data(data)
+    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]
+    iteration_range = iteration_range or (0, 0)
+
+    return XGBPredict(
+        data=data,
+        model=model,
+        output_margin=output_margin,
+        pred_leaf=pred_leaf,
+        pred_contribs=pred_contribs,
+        approx_contribs=approx_contribs,
+        pred_interactions=pred_interactions,
+        validate_features=validate_features,
+        training=training,
+        iteration_range=iteration_range,
+        strict_shape=strict_shape,
+        gpu=data.op.gpu,
+        output_types=output_types,
+        **kwargs,
+    )()

maxframe/learn/contrib/xgboost/regressor.py

@@ -0,0 +1,91 @@
+# 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.
+
+from typing import Union
+
+from ...utils.odpsio import register_odps_model
+from ..utils import make_import_error_func
+from .core import XGBScikitLearnBase, xgboost
+
+if not xgboost:
+    XGBRegressor = make_import_error_func("xgboost")
+else:
+    from xgboost.sklearn import XGBRegressorBase
+
+    from .predict import predict
+
+    @register_odps_model
+    class XGBRegressor(XGBScikitLearnBase, XGBRegressorBase):
+        """
+        Implementation of the scikit-learn API for XGBoost regressor.
+        """
+
+        def __init__(
+            self,
+            xgb_model: Union[xgboost.XGBRegressor, xgboost.Booster] = None,
+            **kwargs,
+        ):
+            super().__init__(**kwargs)
+            self._set_model(xgb_model)
+
+        def fit(
+            self,
+            X,
+            y,
+            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.
+
+            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.
+            """
+            if y.ndim == 2:
+                kw["num_class"] = y.shape[1]
+                kw["output_ndim"] = 2
+            super().fit(
+                X,
+                y,
+                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,
+            )
+            return self
+
+        def predict(self, data, **kw):
+            return predict(self.get_booster(), data, **kw)

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

@@ -0,0 +1,13 @@
+# 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.

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/tests/test_core.py

@@ -0,0 +1,43 @@
+# 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
+
+try:
+    import xgboost
+except ImportError:
+    xgboost = None
+
+
+from ..... import tensor as mt
+
+if xgboost:
+    from ..core import wrap_evaluation_matrices
+
+
+@pytest.mark.skipif(xgboost is None, reason="XGBoost not installed")
+def test_wrap_evaluation_matrices():
+    X = mt.random.rand(100, 3)
+    y = mt.random.randint(3, size=(100,))
+
+    eval_set = [(mt.random.rand(10, 3), mt.random.randint(3, size=10))]
+    with pytest.raises(ValueError):
+        # sample_weight_eval_set size wrong
+        wrap_evaluation_matrices(0.0, X, y, None, None, eval_set, [], None)
+
+    with pytest.raises(ValueError):
+        wrap_evaluation_matrices(0.0, X, y, None, None, None, eval_set, None)
+
+    evals = wrap_evaluation_matrices(0.0, X, y, None, None, eval_set, None, None)[1]
+    assert len(evals) > 0

maxframe/learn/contrib/xgboost/train.py

@@ -0,0 +1,181 @@
+# 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 logging
+from collections import OrderedDict
+from typing import List
+
+from .... import opcodes
+from ....core import EntityData, OutputType
+from ....core.operator import ObjectOperator, ObjectOperatorMixin
+from ....serialization.serializables import (
+    AnyField,
+    BoolField,
+    DictField,
+    FieldTypes,
+    FunctionField,
+    Int16Field,
+    Int64Field,
+    KeyField,
+    ListField,
+)
+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__)
+
+
+def _on_serialize_evals(evals_val):
+    if evals_val is None:
+        return None
+    return [list(x) for x in evals_val]
+
+
+class XGBTrain(ObjectOperator, ObjectOperatorMixin):
+    _op_type_ = opcodes.XGBOOST_TRAIN
+
+    params = DictField("params", key_type=FieldTypes.string, default=None)
+    dtrain = KeyField("dtrain", default=None)
+    evals = ListField("evals", on_serialize=_on_serialize_evals, default=None)
+    obj = FunctionField("obj", default=None)
+    feval = FunctionField("feval", default=None)
+    maximize = BoolField("maximize", default=None)
+    early_stopping_rounds = Int64Field("early_stopping_rounds", default=None)
+    verbose_eval = AnyField("verbose_eval", default=None)
+    xgb_model = KeyField("xgb_model", default=None)
+    callbacks = ListField(
+        "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)
+    num_class = Int64Field("num_class", default=None)
+    _has_evals_result = BoolField("has_evals_result", default=False)
+    output_ndim = Int16Field("output_ndim", default=None)
+
+    def __init__(self, gpu=None, **kw):
+        if kw.get("evals_result") is not None:
+            kw["_has_evals_result"] = True
+        super().__init__(gpu=gpu, **kw)
+        if self.output_types is None:
+            self.output_types = [OutputType.object]
+        if self.has_evals_result:
+            self.output_types.append(OutputType.object)
+
+    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 val in evals_dict.values():
+                new_key = next(input_it)
+                new_evals_dict[new_key] = val
+            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]
+
+    @property
+    def output_limit(self):
+        return 2 if self.has_evals_result else 1
+
+    @property
+    def has_evals_result(self) -> bool:
+        return self._has_evals_result or self.evals
+
+
+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.
+
+    Parameters
+    ----------
+    Parameters are the same as `xgboost.train`. Note that train is an eager-execution
+    API if evals is passed, thus the call will be blocked until training finished.
+
+    Returns
+    -------
+    results: Booster
+    """
+
+    evals_result = evals_result if evals_result is not None else dict()
+    processed_evals = []
+    session = kwargs.pop("session", None)
+    run_kwargs = kwargs.pop("run_kwargs", dict())
+    if evals:
+        for eval_dmatrix, name in evals:
+            if not isinstance(name, str):
+                raise TypeError("evals must a list of pairs (DMatrix, string)")
+            if hasattr(eval_dmatrix, "op") and isinstance(eval_dmatrix.op, ToDMatrix):
+                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)
+    if evals:
+        data.execute(session=session, **run_kwargs)
+    return data