maxframe 2.2.0__cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl

maxframe/learn/contrib/llm/text.py

@@ -0,0 +1,302 @@
+# 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 Any, Dict, List
+
+import numpy as np
+
+from .... import opcodes
+from ....dataframe.core import DataFrame, Series
+from ....serialization.serializables import FieldTypes, ListField, StringField
+from .core import LLM, LLMTaskOperator
+
+
+class TextLLMSummarizeOperator(LLMTaskOperator):
+    _op_type_ = opcodes.LLM_TEXT_SUMMARIZE_TASK
+
+    def get_output_dtypes(self) -> Dict[str, np.dtype]:
+        return {
+            "summary": np.dtype("O"),
+            "success": np.dtype("bool"),
+        }
+
+
+class TextLLMTranslateOperator(LLMTaskOperator):
+    _op_type_ = opcodes.LLM_TEXT_TRANSLATE_TASK
+
+    source_language = StringField("source_language")
+    target_language = StringField("target_language")
+
+    def get_output_dtypes(self) -> Dict[str, np.dtype]:
+        return {
+            "target": np.dtype("O"),
+            "success": np.dtype("bool"),
+        }
+
+
+class TextLLMClassifyOperator(LLMTaskOperator):
+    _op_type_ = opcodes.LLM_TEXT_CLASSIFY_TASK
+
+    labels = ListField("labels")
+    description = StringField("description", default=None)
+    examples = ListField("examples", FieldTypes.dict, default=None)
+
+    def get_output_dtypes(self) -> Dict[str, np.dtype]:
+        return {
+            "label": np.dtype("O"),
+            "reason": np.dtype("O"),
+            "success": np.dtype("bool"),
+        }
+
+
+class TextLLM(LLM):
+    def generate(
+        self,
+        data,
+        prompt_template: List[Dict[str, str]],
+        params: Dict[str, Any] = None,
+    ):
+        raise NotImplementedError
+
+    def summarize(self, series, index=None, **kw):
+        return TextLLMSummarizeOperator(model=self, task="summarize", **kw)(
+            series, index
+        )
+
+    def translate(
+        self,
+        series,
+        target_language: str,
+        source_language: str = None,
+        index=None,
+        **kw
+    ):
+        return TextLLMTranslateOperator(
+            model=self,
+            task="translate",
+            source_language=source_language,
+            target_language=target_language,
+            **kw
+        )(series, index)
+
+    def classify(
+        self,
+        series,
+        labels: List[str],
+        description=None,
+        examples=None,
+        index=None,
+        **kw
+    ):
+        return TextLLMClassifyOperator(
+            model=self,
+            labels=labels,
+            task="classify",
+            description=description,
+            examples=examples,
+            **kw
+        )(series, index)
+
+
+def generate(
+    data,
+    model: TextLLM,
+    prompt_template: List[Dict[str, Any]],
+    params: Dict[str, Any] = None,
+):
+    """
+    Generate text using a text language model based on given data and prompt template.
+
+    Parameters
+    ----------
+    data : DataFrame or Series
+        Input data used for generation. Can be maxframe DataFrame, Series that contain text to be processed.
+    model : TextLLM
+        Language model instance used for text generation.
+    prompt_template : List[Dict[str, str]]
+        Dictionary containing the conversation messages template. Use ``{col_name}`` as a placeholder to reference
+        column data from input data.
+
+        Usually in format of [{"role": "user", "content": "{query}"}], same with openai api schema.
+    params : Dict[str, Any], optional
+        Additional parameters for generation configuration, by default None.
+        Can include settings like temperature, max_tokens, etc.
+
+    Returns
+    -------
+    DataFrame
+        Generated text raw response and success status. If the success is False, the generated text will return the
+        error message.
+
+    Examples
+    --------
+    >>> from maxframe.learn.contrib.llm.models.managed import ManagedTextLLM
+    >>> import maxframe.dataframe as md
+    >>>
+    >>> # Initialize the model
+    >>> llm = ManagedTextLLM(name="Qwen2.5-0.5B-instruct")
+    >>>
+    >>> # Prepare prompt template
+    >>> messages = [
+    ...     {
+    ...         "role": "user",
+    ...         "content": "Help answer following question: {query}",
+    ...     },
+    ... ]
+
+    >>> # Create sample data
+    >>> df = md.DataFrame({"query": ["What is machine learning?"]})
+    >>>
+    >>> # Generate response
+    >>> result = generate(df, llm, prompt_template=messages)
+    >>> result.execute()
+    """
+    if not isinstance(data, DataFrame) and not isinstance(data, Series):
+        raise ValueError("data must be a maxframe dataframe or series object")
+    if not isinstance(model, TextLLM):
+        raise TypeError("model must be a TextLLM object")
+    params = params if params is not None else dict()
+    model.validate_params(params)
+    return model.generate(data, prompt_template=prompt_template, params=params)
+
+
+def summary(series, model: TextLLM, index=None):
+    """
+    Generate summaries for text content in a series using a language model.
+
+    Parameters
+    ----------
+    series : Series
+        A maxframe Series containing text data to be summarized.
+        Each element should be a text string.
+    model : TextLLM
+        Language model instance used for text summarization.
+    index : array-like, optional
+        Index for the output series, by default None, will generate new index.
+
+    Returns
+    -------
+    maxframe.Series
+        A pandas Series containing the generated summaries and success status.
+
+    Notes
+    -----
+      **Preview:** This API is in preview state and may be unstable.
+      The interface may change in future releases.
+    """
+    if not isinstance(series, Series):
+        raise ValueError("series must be a maxframe series object")
+
+    if series.dtype != np.str_:
+        raise ValueError("summary input must be a string series")
+
+    return model.summarize(series, index=index)
+
+
+def translate(
+    series, model: TextLLM, source_language: str, target_language: str, index=None
+):
+    """
+    Translate text content in a series using a language model from source language to target language.
+
+    Parameters
+    ----------
+    series : pandas.Series
+        A maxframe Series containing text data to translate.
+        Each element should be a text string.
+    model : TextLLM
+        Language model instance used for text summarization.
+    source_language : str
+        Source language of the text.
+    target_language : str
+        Target language of the text.
+    index : array-like, optional
+        Index for the output series, by default None, will generate new index.
+
+    Returns
+    -------
+    maxframe.Series
+        A pandas Series containing the generated translation and success status.
+
+    Notes
+    -----
+      **Preview:** This API is in preview state and may be unstable.
+      The interface may change in future releases.
+
+    """
+    if not isinstance(series, Series):
+        raise ValueError("series must be a maxframe series object")
+    if series.dtype != np.str_:
+        raise ValueError("translate input must be a string series")
+    return model.translate(
+        series,
+        source_language=source_language,
+        target_language=target_language,
+        index=index,
+    )
+
+
+def classify(
+    series,
+    model: TextLLM,
+    labels: List[str],
+    description: str = None,
+    examples: List[Dict[str, str]] = None,
+    index=None,
+):
+    """
+    Classify text content in a series with given labels.
+
+    Parameters
+    ----------
+    series : pandas.Series
+        A maxframe Series containing text data to be classified.
+        Each element should be a text string.
+    model : TextLLM
+        Language model instance used for text summarization.
+    labels : List[str]
+        List of labels to classify the text.
+    description : str
+        Description of the classification task.
+    examples : List[Dict[str, Dict[str, str]]]
+        Examples of the classification task, like [{ "text": "text...", "label":"A", reason : "reason..."}], help
+        LLM to better understand your rules.
+    index : array-like, optional
+        Index for the output series, by default None, will generate new index.
+
+    Returns
+    -------
+    maxframe.Series
+        A pandas Series containing the generated classification results and success status.
+
+    Notes
+    -----
+      **Preview:** This API is in preview state and may be unstable.
+      The interface may change in future releases.
+    """
+    if not isinstance(series, Series):
+        raise ValueError("series must be a maxframe series object")
+
+    if series.dtype != np.str_:
+        raise ValueError("classify input must be a string series")
+
+    if not isinstance(labels, list):
+        raise TypeError("labels must be a list")
+
+    if not labels:
+        raise ValueError("labels must not be empty")
+
+    return model.classify(
+        series, labels=labels, description=description, examples=examples, index=index
+    )

maxframe/learn/contrib/models.py

@@ -0,0 +1,106 @@
+# 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 Type
+
+from ... import opcodes
+from ...core import ENTITY_TYPE, OutputType
+from ...core.operator import ObjectOperator, ObjectOperatorMixin
+from ...serialization.serializables import (
+    AnyField,
+    DictField,
+    FunctionField,
+    TupleField,
+)
+from ...udf import BuiltinFunction
+from ...utils import find_objects, replace_objects
+from ..core import Model, ModelData
+
+
+class ModelWithEvalData(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()
+
+    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
+
+
+class ModelWithEval(Model):
+    pass
+
+
+class ModelDataSource(ObjectOperator, ObjectOperatorMixin):
+    _op_type_ = opcodes.MODEL_DATA_SOURCE
+
+    data = AnyField("data")
+
+    def __call__(self, model_cls: Type[ModelWithEval]):
+        self._output_types = [OutputType.object]
+        return self.new_tileable(None, object_class=model_cls)
+
+
+class ModelApplyChunk(ObjectOperator, ObjectOperatorMixin):
+    _op_module_ = "maxframe.learn.contrib.models"
+    _op_type_ = opcodes.APPLY_CHUNK
+
+    func = FunctionField("func")
+    args = TupleField("args")
+    kwargs = DictField("kwargs")
+
+    def __init__(self, output_types=None, **kwargs):
+        if not isinstance(output_types, (tuple, list)):
+            output_types = [output_types]
+        self._output_types = list(output_types)
+        super().__init__(**kwargs)
+
+    def has_custom_code(self) -> bool:
+        return not isinstance(self.func, BuiltinFunction)
+
+    @classmethod
+    def _set_inputs(cls, op: "ModelApplyChunk", inputs):
+        super()._set_inputs(op, inputs)
+        old_inputs = find_objects(op.args, ENTITY_TYPE) + find_objects(
+            op.kwargs, ENTITY_TYPE
+        )
+        mapping = {o: n for o, n in zip(old_inputs, op._inputs[1:])}
+        op.args = replace_objects(op.args, mapping)
+        op.kwargs = replace_objects(op.kwargs, mapping)
+
+    @property
+    def output_limit(self) -> int:
+        return len(self._output_types)
+
+    def __call__(self, t, output_kws, args=None, **kwargs):
+        self.args = args or ()
+        self.kwargs = kwargs
+        inputs = (
+            [t]
+            + find_objects(self.args, ENTITY_TYPE)
+            + find_objects(self.kwargs, ENTITY_TYPE)
+        )
+        return self.new_tileables(inputs, kws=output_kws)
+
+
+def to_remote_model(model, model_cls: Type[ModelWithEval]) -> ModelWithEval:
+    op = ModelDataSource(data=model)
+    return op(model_cls)

maxframe/learn/contrib/pytorch/__init__.py

@@ -0,0 +1,16 @@
+# 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 .run_function import run_pytorch_function
+from .run_script import run_pytorch_script

maxframe/learn/contrib/pytorch/run_function.py

@@ -0,0 +1,110 @@
+# 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 io
+from typing import Any, Callable, Dict
+
+from ....core import TILEABLE_TYPE
+from ....lib import wrapped_pickle as pickle
+from ....typing_ import SessionType
+from ....utils import find_objects, replace_objects
+from .run_script import run_pytorch_script
+
+_script_template = """
+import cloudpickle
+from maxframe.utils import replace_objects
+
+
+def main(**kwargs):
+    func = cloudpickle.loads(%(pickled_func)r)
+    nested = cloudpickle.loads(%(nested)r)
+    args, kw = replace_objects(nested, kwargs)
+    return func(*args, **kw)
+
+
+if __name__ == "__main__":
+    vars = dict()
+    for var_name in cloudpickle.loads(%(var_names)r):
+        vars[var_name] = globals()[var_name]
+    main(**vars)
+"""
+
+
+def run_pytorch_function(
+    func: Callable,
+    args: tuple = (),
+    kwargs: dict = None,
+    *,
+    n_workers: int = 1,
+    retry_when_fail: bool = False,
+    session: SessionType = None,
+    run_kwargs: Dict[str, Any] = None,
+    port: int = None,
+    execute: bool = True,
+):
+    """
+    Run PyTorch function in MaxFrame cluster.
+
+    Besides args and kwargs, the function will receive extra
+    environment variables from the caller:
+
+    * MASTER_ADDR, MASTER_PORT: the endpoint of the master node
+    * RANK: the index of the current worker
+
+    Parameters
+    ----------
+    func: Callable
+        Function or callable object to run
+    args: tuple
+        Args tuple to pass to the function
+    kwargs: dict
+    n_workers : int
+        Number of PyTorch workers
+    retry_when_fail : bool
+        If True, retry when function failed.
+    session
+        MaxFrame session, if not provided, will use default one.
+    run_kwargs : dict
+        Extra kwargs for `session.run`.
+    port : int
+        Port of PyTorch worker or ps, will automatically increase for the same worker
+
+    Returns
+    -------
+    status
+        return {'status': 'ok'} if succeeded, or error raised
+    """
+    kwargs = kwargs or {}
+    packed = [args, kwargs]
+    inputs = find_objects(packed, TILEABLE_TYPE)
+    input_to_var = {t: f"torch_var_{i}" for i, t in enumerate(inputs)}
+    var_to_input = {v: k for k, v in input_to_var.items()}
+    structure = replace_objects(packed, input_to_var)
+
+    replaces = {
+        "pickled_func": pickle.dumps(func),
+        "nested": pickle.dumps(structure),
+        "var_names": pickle.dumps(list(input_to_var.values())),
+    }
+    code = _script_template % replaces
+    return run_pytorch_script(
+        io.StringIO(code),
+        n_workers,
+        data=var_to_input,
+        retry_when_fail=retry_when_fail,
+        session=session,
+        run_kwargs=run_kwargs,
+        port=port,
+        execute=execute,
+    )

maxframe/learn/contrib/pytorch/run_script.py

@@ -0,0 +1,102 @@
+# 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
+import os
+from typing import Any, BinaryIO, Dict, List, Optional, TextIO, Union
+
+from .... import opcodes
+from ....remote.run_script import RunScript, _extract_inputs
+from ....serialization.serializables import Int32Field, StringField
+from ....typing_ import SessionType, TileableType
+from ....utils import to_binary
+
+logger = logging.getLogger(__name__)
+
+
+class RunPyTorch(RunScript):
+    _op_type_ = opcodes.RUN_PYTORCH
+
+    # used for chunk op
+    master_port = Int32Field("master_port", default=None)
+    master_addr = StringField("master_addr", default=None)
+    init_method = StringField("init_method", default=None)
+    master_waiter_name = StringField("master_waiter_name", default=None)
+
+
+def run_pytorch_script(
+    script: Union[bytes, str, BinaryIO, TextIO],
+    n_workers: int,
+    data: Dict[str, TileableType] = None,
+    gpu: Optional[bool] = None,
+    command_argv: List[str] = None,
+    retry_when_fail: bool = False,
+    session: SessionType = None,
+    run_kwargs: Dict[str, Any] = None,
+    port: int = None,
+    execute: bool = True,
+):
+    """
+    Run PyTorch script in MaxFrame cluster.
+
+    Parameters
+    ----------
+    script: str or file-like object
+        Script to run
+    n_workers : int
+        Number of PyTorch workers
+    data : dict
+        Variable name to data.
+    gpu : bool
+        Run PyTorch script on GPU
+    command_argv : list
+        Extra command args for script
+    retry_when_fail : bool
+        If True, retry when function failed.
+    session
+        MaxFrame session, if not provided, will use default one.
+    run_kwargs : dict
+        Extra kwargs for `session.run`.
+    port : int
+        Port of PyTorch worker or ps, will automatically increase for the same worker
+
+    Returns
+    -------
+    status
+        return {'status': 'ok'} if succeeded, or error raised
+    """
+    if int(n_workers) <= 0:
+        raise ValueError("n_workers should be at least 1")
+    if hasattr(script, "read"):
+        code = script.read()
+    else:
+        with open(os.path.abspath(script), "rb") as f:
+            code = f.read()
+
+    inputs = _extract_inputs(data)
+    port = 29500 if port is None else port
+    op = RunPyTorch(
+        data=data,
+        code=to_binary(code),
+        world_size=int(n_workers),
+        retry_when_fail=retry_when_fail,
+        gpu=gpu,
+        master_port=port,
+        command_args=command_argv,
+    )
+    t = op(inputs)
+    if execute:
+        t.execute(session=session, **(run_kwargs or {}))
+    else:
+        return t

maxframe/learn/contrib/pytorch/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/pytorch/tests/test_pytorch.py

@@ -0,0 +1,42 @@
+# 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 ..... import dataframe as md
+from ..... import tensor as mt
+from .....core import TILEABLE_TYPE
+from ..run_function import run_pytorch_function
+
+
+def test_run_function():
+    def test_func(a, b):
+        raise ValueError(a + b)
+
+    df = md.DataFrame(mt.random.rand(10, 10))
+    t = run_pytorch_function(test_func, args=(df, 1), execute=False)
+
+    global_dict = {"__name__": "__main__"}
+    global_dict.update(t.op.data)
+
+    try:
+        exec(t.op.code, global_dict)
+    except ValueError as ex:
+        result = ex.args[0]
+        assert isinstance(result, TILEABLE_TYPE)
+        assert result.op.lhs.key == df.key
+        assert result.op.rhs == 1
+    else:
+        raise AssertionError("Error not raised")
+
+    assert global_dict["torch_var_0"] is df
+    assert "torch_var_1" not in global_dict