maxframe 1.2.0__cp39-cp39-win_amd64.whl → 1.3.0__cp39-cp39-win_amd64.whl

maxframe/learn/contrib/xgboost/classifier.py

@@ -12,9 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Union
+
 import numpy as np
 
-from ....tensor import argmax, transpose
+from .... import tensor as mt
 from ....tensor.merge.vstack import _vstack
 from ..utils import make_import_error_func
 from .core import XGBScikitLearnBase, xgboost
@@ -33,6 +35,14 @@ else:
         Implementation of the scikit-learn API for XGBoost classification.
         """
 
+        def __init__(
+            self,
+            xgb_model: Union[xgboost.XGBClassifier, xgboost.Booster] = None,
+            **kwargs,
+        ):
+            super().__init__(**kwargs)
+            self._set_model(xgb_model)
+
         def fit(
             self,
             X,
@@ -46,7 +56,7 @@ else:
             **kw,
         ):
             session = kw.pop("session", None)
-            run_kwargs = kw.pop("run_kwargs", dict())
+            run_kwargs = kw.pop("run_kwargs", None) or dict()
             dtrain, evals = wrap_evaluation_matrices(
                 None,
                 X,
@@ -58,6 +68,7 @@ else:
                 base_margin_eval_set,
             )
             params = self.get_xgb_params()
+            self._n_features_in = X.shape[1]
             self.n_classes_ = num_class or 1
             if self.n_classes_ > 2:
                 params["objective"] = "multi:softprob"
@@ -81,7 +92,7 @@ else:
         def predict(self, data, **kw):
             prob = self.predict_proba(data, flag=True, **kw)
             if prob.ndim > 1:
-                prediction = argmax(prob, axis=1)
+                prediction = mt.argmax(prob, axis=1)
             else:
                 prediction = (prob > 0.5).astype(np.int64)
             return prediction
@@ -103,7 +114,7 @@ else:
             # binary logistic function
             classone_probs = prediction
             classzero_probs = 1.0 - classone_probs
-            return transpose(_vstack((classzero_probs, classone_probs)))
+            return mt.transpose(_vstack((classzero_probs, classone_probs)))
 
         @property
         def classes_(self) -> np.ndarray:

maxframe/learn/contrib/xgboost/core.py

@@ -12,14 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Callable, Dict, List, Optional, Tuple
+import json
+import os
+import tempfile
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import numpy as np
 
 try:
     import xgboost
 except ImportError:
     xgboost = None
 
+from ....core import OutputType
 from ...core import Model, ModelData
+from ..models import ModelApplyChunk, to_remote_model
 from .dmatrix import DMatrix
 
 
@@ -32,6 +39,33 @@ class BoosterData(ModelData):
         super().__init__(*args, **kwargs)
         self._evals_result = evals_result if evals_result is not None else dict()
 
+    @staticmethod
+    def _get_booster_score(bst, fmap=None, importance_type="weight"):
+        if not fmap:
+            tmp_file_name = ""
+        else:
+            tmp_file = tempfile.NamedTemporaryFile(delete=False)
+            tmp_file.write(fmap)
+            tmp_file.close()
+            tmp_file_name = tmp_file.name
+
+        try:
+            return bst.get_score(fmap=tmp_file_name, importance_type=importance_type)
+        finally:
+            if tmp_file_name:
+                os.unlink(tmp_file_name)
+
+    def get_score(self, fmap="", importance_type="weight"):
+        op = ModelApplyChunk(
+            func=self._get_booster_score, output_types=[OutputType.object]
+        )
+        if not fmap:
+            fmap_data = None
+        else:
+            with open(fmap, "rb") as fmap_file:
+                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)
@@ -82,6 +116,30 @@ else:
         Base class for implementing scikit-learn interface
         """
 
+        def _set_model(
+            self, xgb_model: Union[xgboost.XGBModel, xgboost.Booster] = None
+        ):
+            booster = None
+            if isinstance(xgb_model, xgboost.XGBModel):
+                booster = xgb_model.get_booster()
+            elif isinstance(xgb_model, xgboost.Booster):
+                booster = xgb_model
+
+            if booster is not None:
+                self._Booster = to_remote_model(booster, model_cls=Booster)
+
+        @classmethod
+        def _get_param_names(cls):
+            # make sure `xgb_model` not treated as a model param
+            names = super()._get_param_names()
+            if names:
+                names = [p for p in names if p != "xgb_model"]
+            return names
+
+        def __repr__(self):
+            local_model = self.fetch()
+            return repr(local_model)
+
         def fit(
             self,
             X,
@@ -157,6 +215,55 @@ else:
             self._Booster.execute(session=session, **run_kwargs)
             return super().evals_result()
 
+        def execute(self, session=None, run_kwargs=None):
+            self._Booster.execute(session=session, run_kwargs=run_kwargs)
+            return self
+
+        def fetch(self, session=None, run_kwargs=None):
+            from xgboost import sklearn as xgb_sklearn
+
+            booster = self._Booster.fetch(session=session, run_kwargs=run_kwargs)
+            remote_bst, self._Booster = self._Booster, booster
+            try:
+                local_cls = getattr(xgb_sklearn, type(self).__name__)
+                local_model = local_cls(**self.get_params(deep=True))
+                local_model._Booster = booster
+                return local_model
+            finally:
+                self._Booster = remote_bst
+
+        @staticmethod
+        def _calc_feature_importance(bst, importance_type, n_features):
+            config = json.loads(bst.save_config())
+            bst_type = config["learner"]["gradient_booster"]["name"]
+            dft = "weight" if bst_type == "gblinear" else "gain"
+            importance_type = importance_type or dft
+            score = bst.get_score(importance_type=importance_type)
+            if bst.feature_names is None:
+                feature_names = [f"f{i}" for i in range(n_features)]
+            else:
+                feature_names = bst.feature_names
+            # gblinear returns all features so the `get` in next line is only for gbtree.
+            all_features = [score.get(f, 0.0) for f in feature_names]
+            all_features_arr = np.array(all_features, dtype=np.float32)
+            total = all_features_arr.sum()
+            if total == 0:
+                return all_features_arr
+            return all_features_arr / total
+
+        @property
+        def feature_importances_(self):
+            op = ModelApplyChunk(
+                func=self._calc_feature_importance, output_types=[OutputType.tensor]
+            )
+            params = {"shape": (self._n_features_in,), "dtype": np.dtype(np.float32)}
+            return op(
+                self.get_booster(),
+                [params],
+                importance_type=self.importance_type,
+                n_features=self._n_features_in,
+            )[0]
+
     def wrap_evaluation_matrices(
         missing: float,
         X: Any,

maxframe/learn/contrib/xgboost/dmatrix.py

@@ -141,8 +141,8 @@ def to_dmatrix(
         label_lower_bound=label_lower_bound,
         label_upper_bound=label_upper_bound,
         gpu=data.op.gpu,
-        enable_categorical=enable_categorical,
         _output_types=get_output_types(data),
+        enable_categorical=enable_categorical,
     )
     return op()
 

maxframe/learn/contrib/xgboost/predict.py

@@ -26,7 +26,8 @@ from ....serialization.serializables import (
     TupleField,
 )
 from ....tensor.core import TensorOrder
-from .core import BoosterData
+from ..models import to_remote_model
+from .core import Booster, BoosterData
 from .dmatrix import check_data
 
 
@@ -96,11 +97,15 @@ def predict(
     -------
     results: Booster
     """
+    import xgboost
+
     data = check_data(data)
-    # TODO: check model datatype
+    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(

maxframe/learn/contrib/xgboost/regressor.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Union
 
 from ..utils import make_import_error_func
 from .core import XGBScikitLearnBase, xgboost
@@ -19,15 +20,25 @@ from .core import XGBScikitLearnBase, xgboost
 if not xgboost:
     XGBRegressor = make_import_error_func("xgboost")
 else:
+    from xgboost.sklearn import XGBRegressorBase
+
     from .core import wrap_evaluation_matrices
     from .predict import predict
     from .train import train
 
-    class XGBRegressor(XGBScikitLearnBase):
+    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,
@@ -41,6 +52,9 @@ else:
         ):
             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,

maxframe/learn/contrib/xgboost/train.py

@@ -82,9 +82,8 @@ class XGBTrain(Operator, TileableOperatorMixin):
         inputs = [self.dtrain]
         if self.has_evals_result:
             inputs.extend(e[0] for e in self.evals)
-        return self.new_tileables(
-            inputs, object_class=Booster, evals_result=evals_result
-        )[0]
+        kws = [{"object_class": Booster}, {}]
+        return self.new_tileables(inputs, kws=kws, evals_result=evals_result)[0]
 
     @property
     def output_limit(self):
@@ -129,4 +128,6 @@ def train(params, dtrain, evals=None, evals_result=None, num_class=None, **kwarg
         num_class=num_class,
         **kwargs,
     )(evals_result)
-    return data.execute(session=session, **run_kwargs) if evals else data
+    if evals:
+        data.execute(session=session, **run_kwargs)
+    return data

maxframe/lib/dtypes_extension/__init__.py

@@ -11,4 +11,5 @@
 # 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 .dtypes import ArrowDtype, dict_, infer_arrow_dtype, is_map_dtype
+
+from .dtypes import ArrowDtype, dict_, is_list_dtype, is_map_dtype, list_

maxframe/lib/dtypes_extension/dtypes.py

@@ -11,12 +11,11 @@
 # 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
 
-import numpy as np
 import pandas as pd
 import pyarrow as pa
-from pandas.api.extensions import ExtensionDtype
 
 try:
     from pandas import ArrowDtype
@@ -38,6 +37,15 @@ def dict_(
     return pd.ArrowDtype(pa.map_(key_type, item_type))
 
 
+def list_(value_type: Union[pa.DataType, pa.Field]):
+    """
+    Create ``pd.ArrowDtype(pa.ListType)`` instance from a list or field.
+    """
+    if ArrowDtype is None:
+        raise ImportError("ArrowDtype is not supported in current environment")
+    return pd.ArrowDtype(pa.list_(value_type))
+
+
 def is_map_dtype(dtype: ArrowDtype) -> bool:
     """
     Check whether the dtype is a map type.
@@ -47,45 +55,12 @@ def is_map_dtype(dtype: ArrowDtype) -> bool:
     return isinstance(dtype, ArrowDtype) and isinstance(dtype.pyarrow_dtype, pa.MapType)
 
 
-_dtype_mapping = {
-    pd.Int8Dtype(): lambda x: pa.int8(),
-    pd.Int16Dtype(): lambda x: pa.int16(),
-    pd.Int32Dtype(): lambda x: pa.int32(),
-    pd.Int64Dtype(): lambda x: pa.int64(),
-    pd.UInt8Dtype(): lambda x: pa.uint8(),
-    pd.UInt16Dtype(): lambda x: pa.uint16(),
-    pd.UInt32Dtype(): lambda x: pa.uint32(),
-    pd.UInt64Dtype(): lambda x: pa.uint64(),
-    pd.BooleanDtype(): lambda x: pa.bool_(),
-    pd.Float32Dtype(): lambda x: pa.float32(),
-    pd.Float64Dtype(): lambda x: pa.float64(),
-    pd.StringDtype(): lambda x: pa.string(),
-}
-
-
-def infer_arrow_dtype(
-    dtype: Union[np.dtype, pa.DataType, ExtensionDtype]
-) -> Union[ArrowDtype, ExtensionDtype]:
+def is_list_dtype(dtype: ArrowDtype) -> bool:
     """
-    Convert any pandas accepted dtype to arrow type in a best effort way.
-
-    Parameters
-    ----------
-    dtype : Union[np.dtype, pa.DataType, ExtensionDtype]
-        The dtype instance, can be np.dtype, pa.DataType or ExtensionDtype
-
-    Returns
-    -------
-    Union[pd.ArrowDtype, ExtensionDtype]: The converted pd.ArrowDtype, or ExtensionDtype if conversion failed.
+    Check whether the dtype is a list dtype.
     """
-    if isinstance(dtype, ArrowDtype):
-        return dtype
-
-    if isinstance(dtype, np.dtype):
-        return ArrowDtype(pa.from_numpy_dtype(dtype))
-    if isinstance(dtype, pd.DatetimeTZDtype):
-        return pa.timestamp(dtype.unit, dtype.tz)
-
-    if dtype in _dtype_mapping:
-        return ArrowDtype(_dtype_mapping[dtype](dtype))
-    return dtype
+    if ArrowDtype is None:
+        raise ImportError("ArrowDtype is not supported in current environment")
+    return isinstance(dtype, ArrowDtype) and isinstance(
+        dtype.pyarrow_dtype, pa.ListType
+    )

maxframe/lib/dtypes_extension/tests/test_dtypes.py

@@ -12,24 +12,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import numpy as np
 import pandas as pd
 import pyarrow as pa
 import pytest
 
 from ....utils import ARROW_DTYPE_NOT_SUPPORTED
-from ..dtypes import dict_, infer_arrow_dtype, is_map_dtype
+from ..dtypes import dict_, is_list_dtype, is_map_dtype, list_
 
 try:
     from pandas import ArrowDtype
 except:
     ArrowDtype = None
 
-
-@pytest.mark.skipif(
+pytestmark = pytest.mark.skipif(
     ARROW_DTYPE_NOT_SUPPORTED,
     reason="pandas doesn't support ArrowDtype",
 )
+
+
 def test_map_dtype():
     dt = dict_(pa.int64(), pa.string())
     assert is_map_dtype(dt)
@@ -39,30 +39,10 @@ def test_map_dtype():
     assert not is_map_dtype(pd.Int64Dtype)
 
 
-@pytest.mark.skipif(
-    ARROW_DTYPE_NOT_SUPPORTED,
-    reason="pandas doesn't support ArrowDtype",
-)
-@pytest.mark.parametrize(
-    "input_dtype, expected_type, expected_pa_dtype",
-    [
-        (
-            ArrowDtype(pa.int64()) if ArrowDtype else None,
-            ArrowDtype,
-            pa.int64(),
-        ),  # pd.ArrowDtype
-        (np.dtype("int64"), ArrowDtype, pa.int64()),  # np.dtype
-        (pd.CategoricalDtype(), pd.CategoricalDtype, None),  # pa.DataType
-        (pd.Int64Dtype(), ArrowDtype, pa.int64()),  # pd.ExtensionDtype
-        (
-            pd.DatetimeTZDtype("ns", "Asia/Shanghai"),
-            pa.TimestampType,
-            pa.timestamp("ns", "Asia/Shanghai"),
-        ),
-    ],
-)
-def test_infer_arrow_dtype(input_dtype, expected_type, expected_pa_dtype):
-    result = infer_arrow_dtype(input_dtype)
-    assert isinstance(result, expected_type)
-    if expected_type == ArrowDtype:
-        assert result.pyarrow_dtype == expected_pa_dtype
+def test_list_dtype():
+    dt = list_(pa.int64())
+    assert is_list_dtype(dt)
+
+    dt = pd.ArrowDtype(pa.map_(pa.int64(), pa.string()))
+    assert not is_list_dtype(dt)
+    assert not is_list_dtype(pd.Int64Dtype)

maxframe/opcodes.py

@@ -395,6 +395,11 @@ FUSE = 801
 # LLM
 DASHSCOPE_TEXT_GENERATION = 810
 DASHSCOPE_MULTI_MODAL_GENERATION = 811
+MANAGED_TEXT_MODAL_GENERATION = 812
+MANAGED_MULTI_MODAL_GENERATION = 813
+LLM_TEXT_SUMMARIZE_TASK = 814
+LLM_TEXT_TRANSLATE_TASK = 815
+LLM_TEXT_CLASSIFY_TASK = 816
 
 # table like input for tensor
 TABLE_COO = 1003
@@ -575,16 +580,30 @@ DATAFRAME_RESHUFFLE = 10001
 FLATMAP = 10002
 FLATJSON = 10003
 APPLY_CHUNK = 10004
+
 SERIES_DICT_GETITEM = 10005
 SERIES_DICT_SETITEM = 10006
 SERIES_DICT_LENGTH = 10007
 SERIES_DICT_REMOVE = 10008
 SERIES_DICT_CONTAINS = 10009
+SERIES_DICT_FLATTEN = 10010
+
+SERIES_LIST_GETITEM = 10020
+SERIES_LIST_SETITEM = 10021
+SERIES_LIST_CONTAINS = 10022
+SERIES_LIST_LENGTH = 10023
+SERIES_LIST_INSERT = 10024
+SERIES_LIST_EXTEND = 10025
+SERIES_LIST_POP = 10026
+SERIES_LIST_SORT = 10027
+SERIES_LIST_FLATTEN = 10028
 
 # MaxFrame internal operators
 DATAFRAME_PROJECTION_SAME_INDEX_MERGE = 100001
 GROUPBY_AGGR_SAME_INDEX_MERGE = 100002
 DATAFRAME_ILOC_GET_AND_RENAME_ITEM = 100003
+COLLECT_MODEL_RESULT = 100004
+MODEL_DATA_SOURCE = 100005
 
 # fetches
 FETCH_SHUFFLE = 999998

maxframe/serialization/__init__.py

@@ -16,6 +16,7 @@ from . import arrow, exception, maxframe_objects, numpy, pandas, scipy, serializ
 from .core import (
     PickleContainer,
     Serializer,
+    clear_type_cache,
     deserialize,
     load_type,
     pickle_buffers,

maxframe/serialization/core.pyx

@@ -112,7 +112,14 @@ cpdef object load_type(str class_name, object parent_class):
 
             mod_name, cls_name = class_name.rsplit("#", 1)
 
-            cls = importlib.import_module(mod_name)
+            try:
+                cls = importlib.import_module(mod_name)
+            except ImportError as ex:
+                raise ImportError(
+                    f"Failed to import {mod_name} when loading "
+                    f"class {class_name}, {ex}"
+                ) from None
+
             for sub_cls_name in cls_name.split("."):
                 cls = getattr(cls, sub_cls_name)
         _type_cache[class_name] = cls
@@ -122,6 +129,10 @@ cpdef object load_type(str class_name, object parent_class):
     return cls
 
 
+cpdef void clear_type_cache():
+    _type_cache.clear()
+
+
 cdef Serializer get_deserializer(int32_t deserializer_id):
     return _deserializers[deserializer_id]
 

maxframe/serialization/numpy.py

@@ -70,10 +70,18 @@ class NDArraySerializer(Serializer):
         if dtype_new_order:
             dtype = dtype[dtype_new_order]
         if dtype.hasobject:
-            shape = header["shape"]
-            # fill empty object array
-            val = np.empty(shape, dtype=dtype)
-            val[(slice(None),) * len(shape)] = subs[0]
+            shape = tuple(header["shape"])
+            if shape == ():
+                val = np.array(subs[0]).reshape(shape)
+            else:
+                # fill empty object array
+                val = np.empty(shape, dtype=dtype)
+                try:
+                    val[(slice(None),) * len(shape)] = subs[0]
+                except ValueError:
+                    val[(slice(None),) * len(shape)] = np.array(
+                        subs[0], dtype=dtype
+                    ).reshape(shape)
         else:
             val = np.ndarray(
                 shape=tuple(header["shape"]),

maxframe/serialization/serializables/tests/test_serializable.py

@@ -22,7 +22,7 @@ import pytest
 from ....core import EntityData
 from ....lib.wrapped_pickle import switch_unpickle
 from ....utils import no_default
-from ... import deserialize, serialize
+from ... import clear_type_cache, deserialize, serialize
 from .. import (
     AnyField,
     BoolField,
@@ -202,6 +202,7 @@ def test_serializable(set_is_ci):
 def test_compatible_serializable(set_is_ci):
     global MySimpleSerializable, MySubSerializable
 
+    clear_type_cache()
     old_base, old_sub = MySimpleSerializable, MySubSerializable
 
     try:
@@ -231,13 +232,23 @@ def test_compatible_serializable(set_is_ci):
         my_sub_serializable2 = deserialize(header, buffers)
         assert type(my_sub_serializable) is not type(my_sub_serializable2)
         _assert_serializable_eq(my_sub_serializable, my_sub_serializable2)
+
+        header, buffers = serialize(my_sub_serializable2)
     finally:
         MySimpleSerializable, MySubSerializable = old_base, old_sub
+        MyMidSerializable = None
+        clear_type_cache()
+
+    my_sub_serializable3 = deserialize(header, buffers)
+    assert type(my_sub_serializable2) is not type(my_sub_serializable3)
+    _assert_serializable_eq(my_sub_serializable2, my_sub_serializable3)
 
 
 def _assert_serializable_eq(my_serializable, my_serializable2):
     for field_name, field in my_serializable._FIELDS.items():
-        if not hasattr(my_serializable, field.name):
+        if not hasattr(my_serializable, field.name) or not hasattr(
+            my_serializable2, field.name
+        ):
             continue
         expect_value = getattr(my_serializable, field_name)
         if expect_value is no_default:

maxframe/serialization/tests/test_serial.py

@@ -162,6 +162,8 @@ def test_timezones(val):
 @pytest.mark.parametrize(
     "val",
     [
+        np.array(1234),
+        np.array("str-val", dtype="O"),
         np.array([1024])[0],
         np.array(np.random.rand(100, 100)),
         np.array(np.random.rand(100, 100).T),

maxframe/tensor/merge/concatenate.py

@@ -11,6 +11,7 @@
 # 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 numpy as np
 
 from ... import opcodes