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

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_, is_map_dtype
+
+from .dtypes import ArrowDtype, dict_, is_list_dtype, is_map_dtype, list_

maxframe/lib/dtypes_extension/dtypes.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.
+
 from typing import Union
 
 import pandas as pd
@@ -36,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.
@@ -43,3 +53,14 @@ def is_map_dtype(dtype: ArrowDtype) -> bool:
     if ArrowDtype is None:
         raise ImportError("ArrowDtype is not supported in current environment")
     return isinstance(dtype, ArrowDtype) and isinstance(dtype.pyarrow_dtype, pa.MapType)
+
+
+def is_list_dtype(dtype: ArrowDtype) -> bool:
+    """
+    Check whether the dtype is a list 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

@@ -17,18 +17,19 @@ import pyarrow as pa
 import pytest
 
 from ....utils import ARROW_DTYPE_NOT_SUPPORTED
-from ..dtypes import dict_, 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)
@@ -36,3 +37,12 @@ def test_map_dtype():
     dt = pd.ArrowDtype(pa.list_(pa.int64()))
     assert not is_map_dtype(dt)
     assert not is_map_dtype(pd.Int64Dtype)
+
+
+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

maxframe/tensor/misc/unique.py

@@ -32,23 +32,24 @@ class TensorUnique(TensorHasInput, TensorOperatorMixin):
 
     @property
     def output_limit(self):
-        return 1
+        return 1 + self.return_index + self.return_inverse + self.return_counts
 
-    def _gen_kws(self, input_obj, chunk=False, chunk_index=None):
+    @classmethod
+    def _gen_kws(cls, op: "TensorUnique", input_obj, chunk=False, chunk_index=None):
         kws = []
 
         # unique tensor
         shape = list(input_obj.shape)
-        shape[self.axis] = np.nan
+        shape[op.axis] = np.nan
         kw = {"shape": tuple(shape), "dtype": input_obj.dtype, "gpu": input_obj.op.gpu}
         if chunk:
             idx = [0] * len(shape)
-            idx[self.axis] = chunk_index or 0
+            idx[op.axis] = chunk_index or 0
             kw["index"] = tuple(idx)
         kws.append(kw)
 
         # unique indices tensor
-        if self.return_index:
+        if op.return_index:
             kw = {
                 "shape": (np.nan,),
                 "dtype": np.dtype(np.intp),
@@ -60,9 +61,9 @@ class TensorUnique(TensorHasInput, TensorOperatorMixin):
             kws.append(kw)
 
         # unique inverse tensor
-        if self.return_inverse:
+        if op.return_inverse:
             kw = {
-                "shape": (input_obj.shape[self.axis],),
+                "shape": (input_obj.shape[op.axis],),
                 "dtype": np.dtype(np.intp),
                 "gpu": input_obj.op.gpu,
                 "type": "inverse",
@@ -72,7 +73,7 @@ class TensorUnique(TensorHasInput, TensorOperatorMixin):
             kws.append(kw)
 
         # unique counts tensor
-        if self.return_counts:
+        if op.return_counts:
             kw = {
                 "shape": (np.nan,),
                 "dtype": np.dtype(int),
@@ -92,9 +93,9 @@ class TensorUnique(TensorHasInput, TensorOperatorMixin):
         if self.axis is None:
             if ar.ndim > 1:
                 ar = ar.flatten()
-            self._axis = 0
+            self.axis = 0
         else:
-            self._axis = validate_axis(ar.ndim, self._axis)
+            self.axis = validate_axis(ar.ndim, self.axis)
 
         kws = self._gen_kws(self, ar)
         tensors = self.new_tensors([ar], kws=kws, order=TensorOrder.C_ORDER)

maxframe/tensor/reshape/reshape.py

@@ -181,7 +181,10 @@ def _reshape(a, newshape, order="C", tensor_order=None, out_shape=None):
     if tensor_order is None:
         tensor_order = get_order(order, a.order, available_options="CFA")
     op = TensorReshape(
-        newshape, order, dtype=a.dtype, create_view=tensor_order == a.order
+        newshape=newshape,
+        order=order,
+        dtype=a.dtype,
+        create_view=tensor_order == a.order,
     )
     if out_shape is None:
         out_shape = newshape

maxframe/utils.py

@@ -1140,3 +1140,7 @@ def is_empty(val):
     if isinstance(val, (pd.DataFrame, pd.Series, pd.Index)):
         return val.empty
     return not bool(val)
+
+
+def get_default_table_properties():
+    return {"storagestrategy": "archive"}

{maxframe-1.2.1.dist-info → maxframe-1.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: maxframe
-Version: 1.2.1
+Version: 1.3.0
 Summary: MaxFrame operator-based data analyze framework
 Requires-Dist: numpy<2.0.0,>=1.19.0
 Requires-Dist: pandas>=1.0.0