chore: import upstream snapshot with attribution

swift/__init__.py

@@ -0,0 +1,66 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from typing import TYPE_CHECKING
+
+from .utils.import_utils import _LazyModule
+
+if TYPE_CHECKING:
+    from .agent_template import BaseAgentTemplate, agent_template_map
+    from .arguments import (AppArguments, BaseArguments, DeployArguments, EvalArguments, ExportArguments,
+                            InferArguments, PretrainArguments, RLHFArguments, RolloutArguments, SamplingArguments,
+                            SftArguments)
+    from .callbacks import TrainerCallback, callbacks_map
+    from .dataset import EncodePreprocessor, load_dataset
+    from .infer_engine import (AdapterRequest, GRPOVllmEngine, InferClient, InferEngine, InferRequest, LmdeployEngine,
+                               RequestConfig, SglangEngine, TransformersEngine, VllmEngine)
+    from .loss import BaseLoss, loss_map
+    from .loss_scale import ALL_BASE_STRATEGY, ConfigLossScale, LossScale, get_loss_scale, loss_scale_map
+    from .metrics import InferStats, MeanMetric, eval_metrics_map
+    from .model import get_model_processor, get_processor
+    from .optimizers import OptimizerCallback, optimizers_map
+    from .pipelines import (app_main, deploy_main, eval_main, export_main, infer_main, merge_lora, pretrain_main,
+                            rlhf_main, rollout_main, run_deploy, sampling_main, sft_main)
+    from .template import get_template
+    from .trainers import Seq2SeqTrainer, Seq2SeqTrainingArguments, Trainer, TrainingArguments
+    from .tuner_plugin import PeftTuner, Tuner, tuners_map
+    from .tuners import Swift
+    from .utils import get_logger, safe_snapshot_download
+    from .version import __release_datetime__, __version__
+else:
+    _import_structure = {
+        'version': ['__release_datetime__', '__version__'],
+        'tuners': ['Swift'],
+        'tuner_plugin': ['Tuner', 'PeftTuner', 'tuners_map'],
+        'infer_engine': [
+            'TransformersEngine', 'VllmEngine', 'SglangEngine', 'LmdeployEngine', 'InferRequest', 'RequestConfig',
+            'AdapterRequest', 'InferEngine', 'InferClient', 'GRPOVllmEngine'
+        ],
+        'trainers': ['TrainingArguments', 'Seq2SeqTrainingArguments', 'Trainer', 'Seq2SeqTrainer'],
+        'arguments': [
+            'PretrainArguments', 'SftArguments', 'RLHFArguments', 'ExportArguments', 'InferArguments', 'AppArguments',
+            'EvalArguments', 'SamplingArguments', 'RolloutArguments', 'DeployArguments', 'BaseArguments'
+        ],
+        'pipelines': [
+            'sft_main', 'pretrain_main', 'infer_main', 'rlhf_main', 'export_main', 'app_main', 'eval_main',
+            'sampling_main', 'rollout_main', 'deploy_main', 'merge_lora', 'run_deploy'
+        ],
+        'model': ['get_model_processor', 'get_processor'],
+        'template': ['get_template'],
+        'dataset': ['load_dataset', 'EncodePreprocessor'],
+        'utils': ['get_logger', 'safe_snapshot_download'],
+        'agent_template': ['agent_template_map', 'BaseAgentTemplate'],
+        'loss': ['loss_map', 'BaseLoss'],
+        'metrics': ['eval_metrics_map', 'InferStats', 'MeanMetric'],
+        'optimizers': ['optimizers_map', 'OptimizerCallback'],
+        'callbacks': ['callbacks_map', 'TrainerCallback'],
+        'loss_scale': ['loss_scale_map', 'LossScale', 'get_loss_scale', 'ALL_BASE_STRATEGY', 'ConfigLossScale'],
+    }
+
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

swift/agent_template/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from .base import BaseAgentTemplate
+from .mapping import agent_template_map

swift/agent_template/base.py

@@ -0,0 +1,248 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""
+Agent template module for handling tool calling and function execution.
+
+This module provides base classes and utilities for creating agent templates
+that support tool calling in conversational AI systems. It includes support
+for various agent formats like ReAct, function calling, and parallel execution.
+"""
+import ast
+import json
+from abc import ABC, abstractmethod
+from dataclasses import asdict, dataclass
+from typing import Any, Dict, List, Literal, Optional, Tuple, Union
+
+from swift.infer_engine import Function
+from swift.template import Prompt, split_str_parts_by
+
+
+@dataclass
+class AgentKeyword:
+    action: str = 'Action:'
+    action_input: str = 'Action Input:'
+    observation: str = 'Observation:'
+
+
+@dataclass
+class ToolDesc:
+    name_for_model: str
+    name_for_human: str
+    description_for_model: str
+    parameters: str
+    args_format: str
+
+
+class ReactCompatMixin:
+    """
+    Mixin class providing ReAct-style agent compatibility.
+
+    This mixin handles parsing and formatting of tool calls in the ReAct format,
+    where actions and inputs are marked with specific keywords in the text.
+    """
+    keyword = AgentKeyword()
+
+    @staticmethod
+    def _split_action_action_input(response: str, keyword: AgentKeyword) -> List[Function]:
+        agent_parts = split_str_parts_by(response, list(asdict(keyword).values()))
+        functions = []
+        action_content = None
+
+        for part in agent_parts:
+            key, content = part['key'].lower(), part['content']
+            if action_content is None and key == keyword.action.lower():
+                action_content = content
+            elif action_content is not None and key == keyword.action_input.lower():
+                functions.append(Function(name=action_content, arguments=content))
+                action_content = None
+
+        return functions
+
+    def get_toolcall(self, response: str) -> List[Function]:
+        """
+        Extract tool calls from an agent response.
+
+        Args:
+            response: The agent's response text.
+
+        Returns:
+            List of Function objects representing tool calls.
+        """
+        functions = self._split_action_action_input(response, self.keyword)
+        if len(functions) == 0 and self.keyword != ReactCompatMixin.keyword:
+            # compat react
+            functions = self._split_action_action_input(response, ReactCompatMixin.keyword)
+        return functions
+
+    def _format_tool_responses(
+        self,
+        assistant_content: str,
+        tool_messages,
+    ) -> Tuple[str, 'Prompt']:
+        """
+        Format tool execution results into the conversation.
+
+        Args:
+            assistant_content: The assistant's message containing tool calls.
+            tool_messages: List of tool execution result messages.
+
+        Returns:
+            Tuple of (formatted assistant content, formatted tool responses).
+        """
+        assert len(tool_messages) > 0
+        with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content
+        if with_action:
+            if not assistant_content.endswith(self.keyword.observation):
+                if not assistant_content.endswith('\n'):
+                    assistant_content += '\n'
+                assistant_content += self.keyword.observation
+            res = []
+            for i, tool_message in enumerate(tool_messages):
+                if i > 0:
+                    res.append(self.keyword.observation)
+                tool_content = tool_message['content']
+                res.append(tool_content)
+                if not tool_content.endswith('\n'):
+                    res.append('\n')
+        else:
+            res = []
+            for tool_message in tool_messages:
+                res.append(tool_message['content'])
+        return assistant_content, res
+
+    @staticmethod
+    def _parse_tool_call(content) -> Dict[str, Any]:
+        obj = BaseAgentTemplate._parse_json(content)
+        name = obj['name']
+        arguments = obj.get('arguments')
+        if arguments is None:
+            arguments = obj.get('parameters')
+        arguments = BaseAgentTemplate._parse_json(arguments)
+        assert arguments is not None, f'content: {content}'
+        return {'name': name, 'arguments': arguments}
+
+    def _format_tool_calls(self, tool_call_messages) -> str:
+        """
+        Format tool call messages into ReAct format.
+
+        Args:
+            tool_call_messages: List of messages containing tool call information.
+
+        Returns:
+            Formatted string with Action, Action Input, and Observation markers.
+        """
+        # -> assistant_content
+        tool_calls = []
+        for message in tool_call_messages:
+            tool_call = self._parse_tool_call(message['content'])
+            tool_calls.append(f'{self.keyword.action} {tool_call["name"]}\n'
+                              f'{self.keyword.action_input} {tool_call["arguments"]}\n')
+        tool_calls.append(self.keyword.observation)
+        return ''.join(tool_calls)
+
+
+class BaseAgentTemplate(ReactCompatMixin, ABC):
+    """
+    Abstract base class for agent templates.
+
+    This class provides common functionality for parsing and formatting tools,
+    as well as handling tool calls in different formats. Subclasses must
+    implement the following methods to define their specific behavior:
+    - `_format_tools`: Format tool definitions for the prompt
+    - `_format_tool_calls`: Format tool call messages
+    - `_format_tool_responses`: Format tool execution results
+    - `get_toolcall`: Extract tool calls from agent responses
+    """
+
+    def _add_tool_call_prefix(self, tool_content: str, pre_message=None) -> str:
+        """Hook to prepend a separator before tool_call content based on the preceding message.
+
+        Subclasses can override this to match their jinja template's separator logic
+        (e.g., Qwen3.5/3.6 inserts '\n\n' when assistant has effective content before tool_calls).
+
+        Args:
+            tool_content: The formatted tool_call string from _format_tool_calls.
+            pre_message: The message immediately before the tool_call block, or None.
+
+        Returns:
+            tool_content with any necessary prefix prepended.
+        """
+        return tool_content
+
+    @staticmethod
+    def _get_tool_name(tool):
+        return tool.get('name_for_model') or tool.get('name')
+
+    @staticmethod
+    def unwrap_tool(tool):
+        assert isinstance(tool, dict), f'tool: {tool}'
+        if 'type' in tool and 'function' in tool:
+            tool = tool['function']
+        return tool
+
+    @staticmethod
+    def wrap_tool(tool):
+        assert isinstance(tool, dict), f'tool: {tool}'
+        if 'type' not in tool and 'function' not in tool:
+            tool = {'type': 'function', 'function': tool}
+        return tool
+
+    @staticmethod
+    def _parse_tool(tool, lang: Literal['zh', 'en']) -> ToolDesc:
+        tool = BaseAgentTemplate.unwrap_tool(tool)
+        name_for_model = BaseAgentTemplate._get_tool_name(tool)
+        name_for_human = tool.get('name_for_human') or name_for_model
+
+        description = tool.get('description')
+        if description is None:
+            description = tool.get('description_for_model')
+        parameters = tool.get('parameters') or {}
+        parameters = parameters if isinstance(parameters, str) else json.dumps(parameters, ensure_ascii=False)
+        args_format = '此工具的输入应为JSON对象。' if lang == 'zh' else 'Format the arguments as a JSON object.'
+        tool_desc = ToolDesc(
+            name_for_model=name_for_model,
+            name_for_human=name_for_human,
+            description_for_model=description,
+            parameters=parameters,
+            args_format=args_format)
+        assert name_for_model is not None and description is not None, f'tool_desc: {tool_desc}'
+        return tool_desc
+
+    @staticmethod
+    def _parse_json(json_str: str) -> Optional[Any]:
+        """
+        Parse a JSON string with fallback to ast.literal_eval.
+
+        Args:
+            json_str: String to parse, or already parsed object.
+
+        Returns:
+            Parsed object, or None if parsing fails.
+        """
+        if not isinstance(json_str, str):
+            return json_str
+        try:
+            res = json.loads(json_str)
+        except json.JSONDecodeError:
+            try:
+                res = ast.literal_eval(json_str)
+            except Exception:
+                return
+        return res
+
+    @abstractmethod
+    def _format_tools(self,
+                      tools: List[Union[str, dict]],
+                      system: Optional[str] = None,
+                      user_message: Optional[dict] = None) -> str:
+        """
+        Format tools for inclusion in the agent prompt.
+
+        Args:
+            tools: List of tool definitions (strings or dictionaries).
+            system: System prompt text.
+            user_message: Optional user message to incorporate.
+
+        Returns:
+            Formatted string to include in the prompt.
+        """
+        pass

swift/agent_template/deepseek_v3_1.py

@@ -0,0 +1,87 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import json
+import re
+from typing import List, Optional, Tuple, Union
+
+from swift.infer_engine import Function
+from swift.template import Prompt
+from .base import BaseAgentTemplate
+
+
+class DeepSeekV31AgentTemplate(BaseAgentTemplate):
+
+    def get_toolcall(self, response: str) -> List[Function]:
+        # Parse tool calls using the DSV3.1 format:
+        # <｜tool▁calls▁begin｜><｜tool▁call▁begin｜>name<｜tool▁sep｜>args<｜tool▁call▁end｜>
+        pattern = r'<｜tool▁call▁begin｜>(.*?)<｜tool▁sep｜>(.*?)<｜tool▁call▁end｜>'
+        res_list = re.findall(pattern, response, re.DOTALL)
+        functions = []
+        for name, arguments in res_list:
+            name = name.strip()
+            arguments = self._parse_json(arguments.strip())
+            if arguments is not None:
+                functions.append(Function(name=name, arguments=arguments))
+
+        if len(functions) == 0:
+            # compat react_en
+            return super().get_toolcall(response)
+        return functions
+
+    def _get_tool_responses(self, tool_messages):
+        return ''.join(f'<｜tool▁output▁begin｜>{tool_message["content"]}<｜tool▁output▁end｜>'
+                       for tool_message in tool_messages)
+
+    def _get_tool_calls(self, tool_calls: List[str]):
+        return f'<｜tool▁calls▁begin｜>{"".join(tool_calls)}<｜tool▁calls▁end｜>'
+
+    def _format_tool_responses(
+        self,
+        assistant_content: str,
+        tool_messages,
+    ) -> Tuple[str, 'Prompt']:
+        with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content
+        if with_action:
+            return super()._format_tool_responses(assistant_content, tool_messages)
+        res = ['<｜end▁of▁sentence｜>', self._get_tool_responses(tool_messages)]
+        return assistant_content, res
+
+    def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
+        tool_descs = []
+        system = system or ''
+        for tool in tools:
+            tool = self.unwrap_tool(tool)
+            tool_name = self._get_tool_name(tool)
+            description = tool.get('description', '')
+            parameters = tool.get('parameters', {})
+
+            tool_desc = f"""### {tool_name}
+Description: {description}
+
+Parameters: {json.dumps(parameters, ensure_ascii=False)}"""
+            tool_descs.append(tool_desc)
+
+        tools_section = '\n\n'.join(tool_descs)
+
+        return f"""{system}
+
+## Tools
+You have access to the following tools:
+
+{tools_section}
+
+IMPORTANT: ALWAYS adhere to this exact format for tool use:
+<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>tool_call_name<｜tool▁sep｜>tool_call_arguments<｜tool▁call▁end｜>{{additional_tool_calls}}<｜tool▁calls▁end｜>
+
+Where:
+- `tool_call_name` must be an exact match to one of the available tools
+- `tool_call_arguments` must be valid JSON that strictly follows the tool's Parameters Schema
+- For multiple tool calls, chain them directly without separators or spaces"""
+
+    def _format_tool_calls(self, tool_call_messages):
+        tool_calls = []
+        for message in tool_call_messages:
+            tool_call = self._parse_tool_call(message['content'])
+            name = tool_call['name']
+            arguments = json.dumps(tool_call['arguments'], ensure_ascii=False)
+            tool_calls.append(f'<｜tool▁call▁begin｜>{name}<｜tool▁sep｜>{arguments}<｜tool▁call▁end｜>')
+        return self._get_tool_calls(tool_calls)

swift/agent_template/deepseek_v4.py

@@ -0,0 +1,138 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import json
+import re
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from swift.infer_engine import Function
+from swift.template import Prompt
+from .base import BaseAgentTemplate
+
+DSML_TOKEN = '｜DSML｜'
+
+TOOLS_TEMPLATE = """## Tools
+
+You have access to a set of tools to help answer the user's question. \
+You can invoke tools by writing a "<{dsml_token}tool_calls>" block like the following:
+
+<{dsml_token}tool_calls>
+<{dsml_token}invoke name="$TOOL_NAME">
+<{dsml_token}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</{dsml_token}parameter>
+...
+</{dsml_token}invoke>
+<{dsml_token}invoke name="$TOOL_NAME2">
+...
+</{dsml_token}invoke>
+</{dsml_token}tool_calls>
+
+String parameters should be specified as is and set `string="true"`. \
+For all other types (numbers, booleans, arrays, objects), \
+pass the value in JSON format and set `string="false"`.
+
+If thinking_mode is enabled (triggered by <think>), \
+you MUST output your complete reasoning inside <think>...</think> BEFORE any tool calls or final response.
+
+Otherwise, output directly after </think> with tool calls or final response.
+
+### Available Tool Schemas
+
+{tool_schemas}
+
+You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.
+"""
+
+
+def _to_json(value: Any) -> str:
+    try:
+        return json.dumps(value, ensure_ascii=False)
+    except Exception:
+        return json.dumps(value, ensure_ascii=True)
+
+
+def _encode_arguments_to_dsml(arguments: Dict[str, Any]) -> str:
+    """Encode tool call arguments dict into DSML parameter lines."""
+    lines = []
+    for k, v in arguments.items():
+        is_str = 'true' if isinstance(v, str) else 'false'
+        val = v if isinstance(v, str) else _to_json(v)
+        lines.append(f'<{DSML_TOKEN}parameter name="{k}" string="{is_str}">{val}</{DSML_TOKEN}parameter>')
+    return '\n'.join(lines)
+
+
+class DeepSeekV4AgentTemplate(BaseAgentTemplate):
+
+    def get_toolcall(self, response: str) -> List[Function]:
+        # Parse DSML tool calls from model output
+        # Pattern: <｜DSML｜invoke name="tool_name">...params...</｜DSML｜invoke>
+        invoke_pattern = re.compile(
+            rf'<{re.escape(DSML_TOKEN)}invoke\s+name="([^"]+)">\s*(.*?)\s*</{re.escape(DSML_TOKEN)}invoke>', re.DOTALL)
+        param_pattern = re.compile(
+            rf'<{re.escape(DSML_TOKEN)}parameter\s+name="([^"]+)"\s+string="(true|false)">'
+            rf'(.*?)</{re.escape(DSML_TOKEN)}parameter>', re.DOTALL)
+
+        functions = []
+        for match in invoke_pattern.finditer(response):
+            tool_name = match.group(1)
+            params_block = match.group(2)
+            arguments = {}
+            for pm in param_pattern.finditer(params_block):
+                param_name = pm.group(1)
+                is_string = pm.group(2)
+                param_value = pm.group(3)
+                if is_string == 'false':
+                    try:
+                        param_value = json.loads(param_value)
+                    except json.JSONDecodeError:
+                        pass
+                arguments[param_name] = param_value
+            functions.append(Function(name=tool_name, arguments=json.dumps(arguments, ensure_ascii=False)))
+
+        if len(functions) == 0:
+            # Fallback to ReAct format
+            return super().get_toolcall(response)
+        return functions
+
+    def _get_tool_responses(self, tool_messages):
+        return ''.join(f'<tool_result>{tool_message["content"]}</tool_result>' for tool_message in tool_messages)
+
+    def _format_tool_responses(
+        self,
+        assistant_content: str,
+        tool_messages,
+    ) -> Tuple[str, 'Prompt']:
+        with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content
+        if with_action:
+            return super()._format_tool_responses(assistant_content, tool_messages)
+        res = [
+            '<｜end▁of▁sentence｜><｜User｜>',
+            self._get_tool_responses(tool_messages),
+            '<｜Assistant｜>',
+        ]
+        return assistant_content, res
+
+    def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
+        tool_schemas = []
+        for tool in tools:
+            tool = self.unwrap_tool(tool)
+            tool_schemas.append(_to_json(tool))
+
+        tools_section = TOOLS_TEMPLATE.format(
+            tool_schemas='\n'.join(tool_schemas),
+            dsml_token=DSML_TOKEN,
+        )
+
+        system = system or ''
+        return f'{system}\n\n{tools_section}' if system else tools_section
+
+    def _format_tool_calls(self, tool_call_messages) -> str:
+        invocations = []
+        for message in tool_call_messages:
+            tool_call = self._parse_tool_call(message['content'])
+            name = tool_call['name']
+            arguments = tool_call['arguments']
+            if isinstance(arguments, str):
+                arguments = json.loads(arguments)
+            dsml_args = _encode_arguments_to_dsml(arguments)
+            invocations.append(f'<{DSML_TOKEN}invoke name="{name}">\n{dsml_args}\n</{DSML_TOKEN}invoke>')
+
+        tool_calls_str = '\n'.join(invocations)
+        return f'<{DSML_TOKEN}tool_calls>\n{tool_calls_str}\n</{DSML_TOKEN}tool_calls>'

swift/agent_template/extra.py

@@ -0,0 +1,36 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from typing import List, Optional, Union
+
+from .base import BaseAgentTemplate
+
+
+class ReactGRPOAgentTemplate(BaseAgentTemplate):
+
+    def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
+        tool_names = []
+        tool_descs = []
+        for tool in tools:
+            tool_desc = self._parse_tool(tool, 'en')
+            tool_names.append(tool_desc.name_for_model)
+            tool_descs.append(
+                f'{tool_desc.name_for_model}: Call this tool to interact with the {tool_desc.name_for_human} API. '
+                f'What is the {tool_desc.name_for_human} API useful for? {tool_desc.description_for_model} '
+                f'Parameters: {tool_desc.parameters} {tool_desc.args_format}')
+
+        return """A conversation for tool calling between User and Assistant. The user asks a question which may be solved by calling tools, and the Assistant solves it. The assistant first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning process should be enclosed within <think> </think>tags and answer should follow the ReACT format(Action:xxx\nAction Input:xxx), i.e., <think> reasoning process here </think> Action: action here\nAction Input: parameters here
+
+Answer the following questions as best as you can. You have access to the following tools:
+
+""" + '\n\n'.join(tool_descs) + f"""
+
+Use the following format:
+
+<think>you should always think about what to do</think>
+Action: the action to take, should be one of [{','.join(tool_names)}]
+Action Input: the input to the action
+Observation: the result of the action, given by the actual calling
+... (this Thought/Action/Action Input/Observation can be repeated zero or more times)
+Final Answer: the final answer to the original input question
+
+Begin!
+"""  # noqa

swift/agent_template/gemma4.py

@@ -0,0 +1,210 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import json
+import re
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from swift.infer_engine import Function
+from swift.template import Prompt
+from .base import BaseAgentTemplate
+
+QUOTE = '<|"|>'
+_STANDARD_KEYS = {'description', 'type', 'properties', 'required', 'nullable'}
+
+
+class Gemma4AgentTemplate(BaseAgentTemplate):
+    """Agent template for Google Gemma-4 models.
+
+    Reference: chat_template.jinja shipped with google/gemma-4-12B-it.
+    Tool definitions are wrapped in `<|tool>...<tool|>` and rendered with the
+    custom DSL described by the official chat template.
+    Tool calls follow `<|tool_call>call:NAME{key:value,...}<tool_call|>` and
+    tool responses follow `<|tool_response>response:NAME{...}<tool_response|>`.
+    """
+
+    @classmethod
+    def _format_argument(cls, value: Any, escape_keys: bool = True) -> str:
+        if isinstance(value, bool):
+            return 'true' if value else 'false'
+        if isinstance(value, str):
+            return f'{QUOTE}{value}{QUOTE}'
+        if value is None:
+            return 'null'
+        if isinstance(value, dict):
+            items = []
+            for k in sorted(value.keys()):
+                v = value[k]
+                key_str = f'{QUOTE}{k}{QUOTE}' if escape_keys else str(k)
+                items.append(f'{key_str}:{cls._format_argument(v, escape_keys=escape_keys)}')
+            return '{' + ','.join(items) + '}'
+        if isinstance(value, (list, tuple)):
+            return '[' + ','.join(cls._format_argument(item, escape_keys=escape_keys) for item in value) + ']'
+        return str(value)
+
+    @classmethod
+    def _format_parameters(cls,
+                           properties: Dict[str, Any],
+                           required: Optional[List[str]] = None,
+                           filter_keys: bool = False) -> str:
+        parts = []
+        for key in sorted(properties.keys()):
+            value = properties[key]
+            if filter_keys and key in _STANDARD_KEYS:
+                continue
+            if not isinstance(value, dict):
+                continue
+            inner: List[str] = []
+            type_upper = (value.get('type') or '').upper() if isinstance(value.get('type'), str) else ''
+            if value.get('description'):
+                inner.append(f'description:{QUOTE}{value["description"]}{QUOTE}')
+            if type_upper == 'STRING':
+                if value.get('enum'):
+                    inner.append(f'enum:{cls._format_argument(value["enum"])}')
+            elif type_upper == 'ARRAY':
+                items_value = value.get('items')
+                if isinstance(items_value, dict) and items_value:
+                    items_inner: List[str] = []
+                    items_required = items_value.get('required', [])
+                    for item_key in sorted(items_value.keys()):
+                        item_value = items_value[item_key]
+                        if item_value is None:
+                            continue
+                        if item_key == 'properties' and isinstance(item_value, dict):
+                            items_inner.append(f'properties:{{{cls._format_parameters(item_value, items_required)}}}')
+                        elif item_key == 'required':
+                            req_str = ','.join(f'{QUOTE}{r}{QUOTE}' for r in item_value)
+                            items_inner.append(f'required:[{req_str}]')
+                        elif item_key == 'type':
+                            if isinstance(item_value, str):
+                                items_inner.append(f'type:{cls._format_argument(item_value.upper())}')
+                            else:
+                                items_inner.append(f'type:{cls._format_argument([str(t).upper() for t in item_value])}')
+                        else:
+                            items_inner.append(f'{item_key}:{cls._format_argument(item_value)}')
+                    inner.append('items:{' + ','.join(items_inner) + '}')
+            if value.get('nullable'):
+                inner.append('nullable:true')
+            if type_upper == 'OBJECT':
+                inner_required = value.get('required', [])
+                if isinstance(value.get('properties'), dict):
+                    inner.append(f'properties:{{{cls._format_parameters(value["properties"], inner_required)}}}')
+                else:
+                    inner.append(f'properties:{{{cls._format_parameters(value, inner_required, filter_keys=True)}}}')
+                if value.get('required'):
+                    req_str = ','.join(f'{QUOTE}{r}{QUOTE}' for r in value['required'])
+                    inner.append(f'required:[{req_str}]')
+            inner.append(f'type:{QUOTE}{type_upper}{QUOTE}')
+            parts.append(f'{key}:{{{",".join(inner)}}}')
+        return ','.join(parts)
+
+    @classmethod
+    def _format_function_declaration(cls, tool: Dict[str, Any]) -> str:
+        function = tool['function']
+        name = function.get('name', '')
+        description = function.get('description', '') or ''
+        result = f'declaration:{name}{{description:{QUOTE}{description}{QUOTE}'
+        params = function.get('parameters')
+        if params:
+            param_parts: List[str] = []
+            properties = params.get('properties')
+            if properties:
+                param_parts.append(f'properties:{{{cls._format_parameters(properties, params.get("required", []))}}}')
+            if params.get('required'):
+                req_str = ','.join(f'{QUOTE}{r}{QUOTE}' for r in params['required'])
+                param_parts.append(f'required:[{req_str}]')
+            ptype = params.get('type')
+            if isinstance(ptype, str) and ptype:
+                param_parts.append(f'type:{QUOTE}{ptype.upper()}{QUOTE}')
+            if param_parts:
+                result += ',parameters:{' + ','.join(param_parts) + '}'
+        result += '}'
+        return result
+
+    def _format_tools(self,
+                      tools: List[Union[str, dict]],
+                      system: Optional[str] = None,
+                      user_message: Optional[dict] = None) -> str:
+        tool_blocks: List[str] = []
+        for tool in tools:
+            tool = self.wrap_tool(tool)
+            tool_blocks.append(f'<|tool>{self._format_function_declaration(tool)}<tool|>')
+        system_text = (system or '').strip()
+        return system_text + ''.join(tool_blocks)
+
+    def _format_tool_calls(self, tool_call_messages) -> str:
+        invocations: List[str] = []
+        for message in tool_call_messages:
+            tool_call = self._parse_tool_call(message['content'])
+            name = tool_call['name']
+            arguments = tool_call['arguments']
+            if isinstance(arguments, str):
+                arguments = self._parse_json(arguments) or {}
+            if isinstance(arguments, dict):
+                args_str = ','.join(f'{k}:{self._format_argument(arguments[k], escape_keys=False)}'
+                                    for k in sorted(arguments.keys()))
+            else:
+                args_str = ''
+            invocations.append(f'<|tool_call>call:{name}{{{args_str}}}<tool_call|>')
+        return ''.join(invocations)
+
+    def _get_tool_responses(self, tool_messages) -> str:
+        parts: List[str] = []
+        for tool_message in tool_messages:
+            tool_name = tool_message.get('name') or 'unknown'
+            tool_content = tool_message.get('content')
+            if isinstance(tool_content, dict):
+                inner = ','.join(f'{k}:{self._format_argument(tool_content[k], escape_keys=False)}'
+                                 for k in sorted(tool_content.keys()))
+                parts.append(f'<|tool_response>response:{tool_name}{{{inner}}}<tool_response|>')
+            else:
+                # Match jinja: treat string/other content as a single `value:` field.
+                value = '' if tool_content is None else tool_content
+                parts.append(f'<|tool_response>response:{tool_name}'
+                             f'{{value:{self._format_argument(value, escape_keys=False)}}}<tool_response|>')
+        return ''.join(parts)
+
+    def _format_tool_responses(
+        self,
+        assistant_content: str,
+        tool_messages,
+    ) -> Tuple[str, 'Prompt']:
+        with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content
+        if with_action:
+            return super()._format_tool_responses(assistant_content, tool_messages)
+        # If the model hallucinated a trailing `<|tool_response>` opener (e.g. when stop
+        # tokens were not configured), strip it so the rendered turn does not contain
+        # `<|tool_response><|tool_response>response:...`.
+        if assistant_content.endswith('<|tool_response>'):
+            assistant_content = assistant_content[:-len('<|tool_response>')]
+        # In gemma4, tool_call/tool_response/follow-up assistant text all live in the
+        # same `<|turn>model ... <turn|>` block, so we do not open a new model turn here.
+        res: 'Prompt' = [self._get_tool_responses(tool_messages)]
+        return assistant_content, res
+
+    @classmethod
+    def _gemma_to_json(cls, s: str) -> str:
+        # `<|"|>` -> `"`; bare keys preceded by `{` or `,` get JSON-quoted.
+        s = s.replace(QUOTE, '"')
+        s = re.sub(r'(?<=[\{,])([A-Za-z_][\w\-]*)(?=:)', r'"\1"', s)
+        return s
+
+    @classmethod
+    def _parse_arguments(cls, args_body: str) -> Dict[str, Any]:
+        json_str = cls._gemma_to_json('{' + args_body + '}')
+        try:
+            parsed = json.loads(json_str)
+            if isinstance(parsed, dict):
+                return parsed
+        except json.JSONDecodeError:
+            pass
+        return {}
+
+    def get_toolcall(self, response: str) -> List[Function]:
+        pattern = re.compile(r'<\|tool_call>call:([^\{]+)\{(.*?)\}<tool_call\|>', re.DOTALL)
+        functions: List[Function] = []
+        for match in pattern.finditer(response):
+            name = match.group(1).strip()
+            arguments = self._parse_arguments(match.group(2))
+            functions.append(Function(name=name, arguments=json.dumps(arguments, ensure_ascii=False)))
+        if not functions:
+            return super().get_toolcall(response)
+        return functions

swift/agent_template/glm4.py

@@ -0,0 +1,178 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import json
+import re
+from typing import List, Optional, Tuple, Union
+
+from swift.infer_engine import Function
+from swift.template import Prompt
+from .base import BaseAgentTemplate
+
+
+class ChatGLM4AgentTemplate(BaseAgentTemplate):
+    is_glm4_0414 = False
+
+    @staticmethod
+    def _find_function_call(single_content: str) -> Optional[Function]:
+        single_content = single_content.replace('<|observation|>', '')
+        pattern = re.compile(r'([^\n`]*?)\n({.*?})(?=\w*\n|$)', re.DOTALL)
+        matches = pattern.findall(single_content)
+        if not matches:
+            return
+        name, arguments = matches[0]
+        return Function(name=name, arguments=arguments)
+
+    def get_toolcall(self, response: str) -> List[Function]:
+        toolcall_list = response.split('<|assistant|>')
+        functions = []
+        for toolcall in toolcall_list:
+            function = self._find_function_call(toolcall)
+            if function:
+                functions.append(function)
+        if len(functions) == 0:
+            # compat react_en
+            return super().get_toolcall(response)
+        return functions
+
+    def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
+        tool_descs = []
+        for tool in tools:
+            tool = self.unwrap_tool(tool)
+            name = self._get_tool_name(tool)
+            tool_descs.append(f'## {name}\n\n{json.dumps(tool, ensure_ascii=False, indent=4)}\n'
+                              '在调用上述函数时，请使用 Json 格式表示调用的参数。')
+        glm4_system = '你是一个名为 GLM-4 的人工智能助手。你是基于智谱AI训练的语言模型 GLM-4 模型开发的，你的任务是针对用户的问题和要求提供适当的答复和支持。\n\n'  # noqa
+        return ('' if self.is_glm4_0414 else glm4_system) + """# 可用工具
+
+""" + '\n'.join(tool_descs)
+
+    def _format_tool_responses(
+        self,
+        assistant_content: str,
+        tool_messages,
+    ) -> Tuple[str, 'Prompt']:
+        with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content
+        if with_action:
+            return super()._format_tool_responses(assistant_content, tool_messages)
+        res = ['\n']
+        for i, tool_message in enumerate(tool_messages):
+            tool_content = tool_message['content']
+            if i > 0:
+                res.append('<|observation|>\n')
+            res.append(tool_content)
+        res.append('<|assistant|>\n')
+        return assistant_content, res
+
+    def _format_tool_calls(self, tool_call_messages) -> str:
+        tool_calls = []
+        for message in tool_call_messages:
+            tool_call = self._parse_tool_call(message['content'])
+            tool_calls.append(f'{tool_call["name"]}\n{tool_call["arguments"]}')
+        return '<|assistant|>'.join(tool_calls) + '<|observation|>'
+
+
+class GLM4AgentTemplate(ChatGLM4AgentTemplate):
+    is_glm4_0414 = True
+
+
+class GLM4_5AgentTemplate(BaseAgentTemplate):
+    model_type = 'glm4_5'
+
+    @staticmethod
+    def _find_function_call(single_content: str) -> Optional[Function]:
+        single_content = single_content.strip()
+        func_name_match = re.match(r'^([^\n<]+)', single_content)
+        if not func_name_match:
+            return None
+        func_name = func_name_match.group(1).strip()
+        keys = re.findall(r'<arg_key>(.*?)</arg_key>', single_content, re.DOTALL)
+        values = re.findall(r'<arg_value>(.*?)</arg_value>', single_content, re.DOTALL)
+        if len(keys) != len(values):
+            return None
+        args = {k.strip(): v.strip() for k, v in zip(keys, values)}
+        return Function(name=func_name, arguments=json.dumps(args, ensure_ascii=False))
+
+    def get_toolcall(self, response: str) -> List[Function]:
+        toolcall_list = re.findall(r'<tool_call>(.*?)</tool_call>', response, re.DOTALL)
+        functions = []
+        for toolcall in toolcall_list:
+            function = self._find_function_call(toolcall)
+            if function:
+                functions.append(function)
+        if len(functions) == 0:
+            # compat react_en
+            return super().get_toolcall(response)
+        return functions
+
+    def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
+        tool_descs = [
+            '# Tools\n\nYou may call one or more functions to assist with the user query.\n\n'
+            'You are provided with function signatures within <tools></tools> XML tags:\n<tools>'
+        ]
+        for tool in tools:
+            if self.model_type == 'glm5_1':
+                tool = self.unwrap_tool(tool)
+            tool_descs.append(f'{json.dumps(tool, ensure_ascii=False)}')
+        if self.model_type == 'glm4_5':
+            tool_desc = ('</tools>\n\nFor each function call, output the function name and arguments within '
+                         'the following XML format:\n<tool_call>{function-name}\n<arg_key>{arg-key-1}</arg_key>\n'
+                         '<arg_value>{arg-value-1}</arg_value>\n<arg_key>{arg-key-2}</arg_key>\n'
+                         '<arg_value>{arg-value-2}</arg_value>\n...\n</tool_call>')
+        elif self.model_type in {'glm4_7', 'glm5_1'}:
+            tool_desc = ('</tools>\n\nFor each function call, output the function name and arguments within '
+                         'the following XML format:\n<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key>'
+                         '<arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>'
+                         '{arg-value-2}</arg_value>...</tool_call>')
+        else:
+            raise ValueError("model_type must be one of 'glm4_5', 'glm4_7', or 'glm5_1'.")
+        tool_descs.append(tool_desc)
+        tool_descs = '\n'.join(tool_descs)
+        if system is not None and system.strip():
+            tool_descs += '<|system|>\n' + system.strip()
+        elif self.model_type in {'glm4_7', 'glm5_1'} and not tool_descs.startswith('\n'):
+            tool_descs = '\n' + tool_descs
+        return tool_descs
+
+    def _format_tool_responses(
+        self,
+        assistant_content: str,
+        tool_messages,
+    ) -> Tuple[str, 'Prompt']:
+        with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content
+        if with_action:
+            return super()._format_tool_responses(assistant_content, tool_messages)
+        if self.model_type == 'glm4_5':
+            res = []
+            for tool_message in tool_messages:
+                tool_content = tool_message['content']
+                res.append(f'\n<tool_response>\n{tool_content}\n</tool_response>')
+            res.append('<|assistant|>\n')
+        elif self.model_type in {'glm4_7', 'glm5_1'}:
+            res = []
+            for tool_message in tool_messages:
+                tool_content = tool_message['content']
+                res.append(f'<tool_response>{tool_content}</tool_response>')
+            res.append('<|assistant|>')
+        return assistant_content, res
+
+    def _format_tool_calls(self, tool_call_messages) -> str:
+        tool_calls = []
+        for message in tool_call_messages:
+            tool_call = self._parse_tool_call(message['content'])
+            tool_calls.append(f"<tool_call>{tool_call['name']}")
+            for arg_key, arg_value in tool_call['arguments'].items():
+                tool_calls.append(f'<arg_key>{arg_key}</arg_key>')
+                tool_calls.append(f'<arg_value>{arg_value}</arg_value>')
+            tool_calls.append('</tool_call>')
+        if self.model_type == 'glm4_5':
+            sep = '\n'
+        elif self.model_type in {'glm4_7', 'glm5_1'}:
+            sep = ''
+        return sep.join(tool_calls) + '<|observation|>'
+
+
+class GLM4_7AgentTemplate(GLM4_5AgentTemplate):
+    model_type = 'glm4_7'
+
+
+class GLM5_1AgentTemplate(GLM4_5AgentTemplate):
+    model_type = 'glm5_1'

swift/agent_template/hermes.py

@@ -0,0 +1,127 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import json
+import re
+from typing import List, Optional, Tuple, Union
+
+from swift.infer_engine import Function
+from swift.template import Prompt
+from .base import BaseAgentTemplate
+
+
+class HermesAgentTemplate(BaseAgentTemplate):
+
+    def get_toolcall(self, response: str) -> List[Function]:
+        res_list = re.findall(r'<tool_call>(.+?)</tool_call>', response, re.DOTALL)
+        functions = []
+        for res in res_list:
+            res = self._parse_json(res)
+            if isinstance(res, dict) and 'name' in res and 'arguments' in res:
+                functions.append(Function(name=res['name'], arguments=res['arguments']))
+        if len(functions) == 0:
+            # compat react_en
+            return super().get_toolcall(response)
+        return functions
+
+    def _get_tool_responses(self, tool_messages):
+        res_tool = []
+        for tool_message in tool_messages:
+            tool_content = tool_message['content']
+            res_tool.append(f'<tool_response>\n{tool_content}\n</tool_response>')
+        return '\n'.join(res_tool)
+
+    def _get_tool_calls(self, tool_calls: List[str]):
+        return '\n'.join(tool_calls)
+
+    def _format_tool_responses(
+        self,
+        assistant_content: str,
+        tool_messages,
+    ) -> Tuple[str, 'Prompt']:
+        with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content
+        if with_action:
+            return super()._format_tool_responses(assistant_content, tool_messages)
+        if hasattr(self, 'template_meta'):
+            prompt = self.template_meta.prompt
+            chat_sep = self.template_meta.chat_sep
+        else:
+            prompt = ['<|im_start|>user\n{{QUERY}}<|im_end|>\n<|im_start|>assistant\n']
+            chat_sep = ['<|im_end|>\n']
+        res = chat_sep.copy()
+        total_tool = self._get_tool_responses(tool_messages)
+        for context in prompt:
+            if isinstance(context, str):
+                context = context.replace('{{QUERY}}', total_tool)
+            res.append(context)
+        return assistant_content, res
+
+    def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
+        tool_descs = [json.dumps(self.wrap_tool(tool), ensure_ascii=False) for tool in tools]
+        system = system or ''
+        return f"""{system}
+
+# Tools
+
+You may call one or more functions to assist with the user query.
+
+You are provided with function signatures within <tools></tools> XML tags:
+<tools>
+""" + '\n'.join(tool_descs) + """
+</tools>
+
+For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:
+<tool_call>
+{"name": <function-name>, "arguments": <args-json-object>}
+</tool_call>"""
+
+    def _format_tool_calls(self, tool_call_messages):
+        tool_calls = []
+        for message in tool_call_messages:
+            tool_call = self._parse_tool_call(message['content'])
+            tool_calls.append(f'<tool_call>\n{json.dumps(tool_call, ensure_ascii=False)}\n</tool_call>')
+        return self._get_tool_calls(tool_calls)
+
+
+class HunyuanHermesAgentTemplate(HermesAgentTemplate):
+
+    def get_toolcall(self, response: str) -> List[Function]:
+        res_list = re.findall(r'<tool_call>(.+?)\n```json(.+?)```</tool_call>', response, re.DOTALL)
+        functions = []
+        for name, arguments in res_list:
+            arguments = self._parse_json(arguments)
+            functions.append(Function(name=name, arguments=arguments))
+        if len(functions) == 0:
+            # compat react_en
+            return super().get_toolcall(response)
+        return functions
+
+    def _get_tool_responses(self, tool_messages):
+        res_tool = []
+        for tool_message in tool_messages:
+            tool_content = tool_message['content']
+            res_tool.append(f'<tool_response>{tool_content}</tool_response>')
+        tool_responses = '\n'.join(res_tool)
+        return f'<tool_responses>{tool_responses}</tool_responses>'
+
+    def _get_tool_calls(self, tool_calls: List[str]):
+        tool_calls = '\n'.join(tool_calls)
+        return f'<tool_calls>\n{tool_calls}\n</tool_calls>'
+
+    def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
+        tool_descs = [json.dumps(self.wrap_tool(tool), ensure_ascii=False) for tool in tools]
+        system = system or ''
+        if system:
+            system = f'{system}\n\n'
+        return f"""{system}# Tools
+
+You may call one or more functions to assist with the user query.
+
+You are provided with function signatures within <tools></tools> XML tags:
+<tools>
+""" + '\n'.join(tool_descs) + """
+</tools>
+
+For function call returns, you should first print <tool_calls>For each function call, you should return object like:
+<tool_call>function_name
+```json
+function_arguments_in_json_format
+```</tool_call>At the end of function call returns, you should print </tool_calls>"""

swift/agent_template/hy_v3.py

@@ -0,0 +1,107 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import json
+import re
+from typing import List, Optional, Tuple, Union
+
+from swift.infer_engine import Function
+from swift.template import Prompt
+from .base import BaseAgentTemplate
+
+
+class HyV3PreviewAgentTemplate(BaseAgentTemplate):
+    HYTK = ''
+
+    def get_toolcall(self, response: str) -> List[Function]:
+        # Parse tool calls from <tool_calls>...<tool_call>name<tool_sep>...<arg_key>...<arg_value>...</tool_call>...
+        tool_call_blocks = re.findall(rf'<tool_call{self.HYTK}>(.*?)</tool_call{self.HYTK}>', response, re.DOTALL)
+        functions = []
+        for block in tool_call_blocks:
+            # Extract function name: text before <tool_sep>
+            name_match = re.match(rf'(.*?)<tool_sep{self.HYTK}>', block, re.DOTALL)
+            if not name_match:
+                continue
+            name = name_match.group(1).strip()
+            # Extract arg_key/arg_value pairs together to avoid misalignment
+            pairs = re.findall(
+                rf'<arg_key{self.HYTK}>(.*?)</arg_key{self.HYTK}>\s*<arg_value{self.HYTK}>(.*?)</arg_value{self.HYTK}>',
+                block, re.DOTALL)
+            arguments = {}
+            for k, v in pairs:
+                k = k.strip()
+                v = v.strip()
+                parsed = self._parse_json(v)
+                arguments[k] = parsed if parsed is not None else v
+            functions.append(Function(name=name, arguments=arguments))
+        if len(functions) == 0:
+            # compat react_en
+            return super().get_toolcall(response)
+        return functions
+
+    def _get_tool_responses(self, tool_messages):
+        res_tool = []
+        for tool_message in tool_messages:
+            tool_content = tool_message['content']
+            res_tool.append(f'<tool_response{self.HYTK}>\n{tool_content}\n</tool_response{self.HYTK}>')
+        tool_responses = '\n'.join(res_tool)
+        return f'<tool_responses{self.HYTK}>\n{tool_responses}\n</tool_responses{self.HYTK}>'
+
+    def _get_tool_calls(self, tool_calls: List[str]):
+        tool_calls_str = '\n'.join(tool_calls)
+        return f'<tool_calls{self.HYTK}>\n{tool_calls_str}\n</tool_calls{self.HYTK}>'
+
+    def _format_tool_responses(
+        self,
+        assistant_content: str,
+        tool_messages,
+    ) -> Tuple[str, 'Prompt']:
+        with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content
+        if with_action:
+            return super()._format_tool_responses(assistant_content, tool_messages)
+        res = [f'<｜hy_eos{self.HYTK}｜>', self._get_tool_responses(tool_messages), f'<｜hy_Assistant{self.HYTK}｜>']
+        return assistant_content, res
+
+    def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
+        tool_descs = [json.dumps(self.wrap_tool(tool), ensure_ascii=False) for tool in tools]
+        system = system or ''
+        if system:
+            system = f'{system}\n\n'
+        return f"""{system}# Tools
+
+You may call one or more functions to assist with the user query.
+
+You are provided with function signatures within <tools></tools> XML tags:
+<tools>
+""" + '\n'.join(tool_descs) + f"""
+</tools>
+
+For function call returns, you should first print <tool_calls{self.HYTK}>
+For each function call, you should return object like:
+<tool_call{self.HYTK}>{{function-name}}<tool_sep{self.HYTK}>
+<arg_key{self.HYTK}>{{arg-key-1}}</arg_key{self.HYTK}>
+<arg_value{self.HYTK}>{{arg-value-1}}</arg_value{self.HYTK}>
+<arg_key{self.HYTK}>{{arg-key-2}}</arg_key{self.HYTK}>
+<arg_value{self.HYTK}>{{arg-value-2}}</arg_value{self.HYTK}>
+...
+</tool_call{self.HYTK}>
+At the end of function call returns, you should print </tool_calls{self.HYTK}>"""
+
+    def _format_tool_calls(self, tool_call_messages):
+        tool_calls = []
+        for message in tool_call_messages:
+            tool_call = self._parse_tool_call(message['content'])
+            name = tool_call['name']
+            arguments = tool_call['arguments']
+            arg_lines = []
+            if isinstance(arguments, dict):
+                for k, v in arguments.items():
+                    if not isinstance(v, str):
+                        v = json.dumps(v, ensure_ascii=False)
+                    arg_lines.append(f'<arg_key{self.HYTK}>{k}</arg_key{self.HYTK}>\n'
+                                     f'<arg_value{self.HYTK}>{v}</arg_value{self.HYTK}>')
+            arg_str = '\n'.join(arg_lines)
+            tool_calls.append(f'<tool_call{self.HYTK}>{name}<tool_sep{self.HYTK}>\n{arg_str}\n</tool_call{self.HYTK}>')
+        return self._get_tool_calls(tool_calls)
+
+
+class HyV3AgentTemplate(HyV3PreviewAgentTemplate):
+    HYTK = ':opensource'

swift/agent_template/kimi_k25.py

@@ -0,0 +1,119 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import json
+import re
+from typing import List, Optional, Tuple, Union
+
+from swift.infer_engine import Function
+from swift.template import Prompt
+from .base import BaseAgentTemplate
+
+
+class KimiK25AgentTemplate(BaseAgentTemplate):
+    """Agent template for Kimi K2.5/K2.6 models.
+
+    Tool calling format:
+    - Tools are declared in a separate system message with role 'tool_declare'
+      using TypeScript namespace format.
+    - Tool calls:
+      <|tool_calls_section_begin|>
+      <|tool_call_begin|>{function_name}<|tool_call_argument_begin|>{args_json}<|tool_call_end|>
+      <|tool_calls_section_end|>
+    - Tool response:
+      <|im_system|>tool<|im_middle|>## Return of {tool_call_id}
+      {content}<|im_end|>
+    """
+
+    @staticmethod
+    def _json_type_to_ts(json_type):
+        type_map = {
+            'string': 'string',
+            'number': 'number',
+            'integer': 'number',
+            'boolean': 'boolean',
+            'array': 'any[]',
+            'object': 'object',
+            'null': 'null',
+        }
+        return type_map.get(json_type, 'any')
+
+    def _tools_to_typescript(self, tools):
+        parts = []
+        for tool in tools:
+            tool = self.unwrap_tool(tool)
+            name = self._get_tool_name(tool)
+            description = tool.get('description', '')
+            parameters = tool.get('parameters', {})
+            properties = parameters.get('properties', {})
+
+            lines = []
+            if description:
+                lines.append(f'// {description}')
+
+            if not properties:
+                lines.append(f'type {name} = (_: {{}}) => any;')
+            else:
+                lines.append(f'type {name} = (_: {{')
+                props = list(properties.items())
+                for i, (pname, pschema) in enumerate(props):
+                    pdesc = pschema.get('description', '')
+                    ptype = self._json_type_to_ts(pschema.get('type', ''))
+                    if pdesc:
+                        lines.append(f'  // {pdesc}')
+                    if i < len(props) - 1:
+                        lines.append(f'  {pname}: {ptype},')
+                    else:
+                        lines.append(f'  {pname}: {ptype}')
+                lines.append('}) => any;')
+            parts.append('\n'.join(lines))
+        return '\n'.join(parts)
+
+    def get_toolcall(self, response: str) -> List[Function]:
+        pattern = r'<\|tool_call_begin\|>(.*?)<\|tool_call_argument_begin\|>(.*?)<\|tool_call_end\|>'
+        res_list = re.findall(pattern, response, re.DOTALL)
+        functions = []
+        for name, arguments in res_list:
+            name = name.strip()
+            arguments = arguments.strip()
+            parsed_args = self._parse_json(arguments)
+            if parsed_args is not None:
+                functions.append(Function(name=name, arguments=parsed_args))
+            else:
+                functions.append(Function(name=name, arguments=arguments))
+        if len(functions) == 0:
+            # compat react_en
+            return super().get_toolcall(response)
+        return functions
+
+    def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
+        ts_tools = self._tools_to_typescript(tools)
+        tool_content = f'# Tools\n\n## functions\nnamespace functions {{\n{ts_tools}\n}}\n'
+        system = system or ''
+        res = f'tool_declare<|im_middle|>{tool_content}'
+        if system:
+            res += f'<|im_end|><|im_system|>system<|im_middle|>{system}'
+        return res
+
+    def _format_tool_calls(self, tool_call_messages) -> str:
+        tool_calls = []
+        for message in tool_call_messages:
+            tool_call = self._parse_tool_call(message['content'])
+            name = tool_call['name']
+            arguments = json.dumps(tool_call['arguments'], ensure_ascii=False)
+            tool_calls.append(f'<|tool_call_begin|>{name}<|tool_call_argument_begin|>{arguments}<|tool_call_end|>')
+        return f'<|tool_calls_section_begin|>{"".join(tool_calls)}<|tool_calls_section_end|>'
+
+    def _format_tool_responses(
+        self,
+        assistant_content: str,
+        tool_messages,
+    ) -> Tuple[str, 'Prompt']:
+        with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content
+        if with_action:
+            return super()._format_tool_responses(assistant_content, tool_messages)
+        res = ['<|im_end|>']
+        for tool_message in tool_messages:
+            tool_call_id = tool_message.get('tool_call_id', '')
+            tool_content = tool_message['content']
+            res.append(f'<|im_system|>tool<|im_middle|>## Return of {tool_call_id}\n{tool_content}<|im_end|>')
+        res.append('<|im_assistant|>assistant<|im_middle|>')
+        return assistant_content, res

swift/agent_template/llama.py

@@ -0,0 +1,74 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import json
+import re
+from typing import List, Optional, Tuple, Union
+
+from swift.infer_engine import Function
+from swift.template import Prompt
+from .base import BaseAgentTemplate
+
+
+class Llama3AgentTemplate(BaseAgentTemplate):
+    eom_token = '<|eom_id|>'
+    start_token = '<|start_header_id|>'
+    end_token = '<|end_header_id|>'
+    eot_token = '<|eot_id|>'
+
+    def get_toolcall(self, response: str) -> List[Function]:
+        if response.endswith(self.eom_token):
+            response = response[:-len(self.eom_token)]
+        functions = []
+        res_list = re.findall(r'{[^{]*?"name":.*?"parameters":\s*?{.*?}\s*?}', response, re.DOTALL)
+        for res in res_list:
+            res = self._parse_json(res)
+            if isinstance(res, dict) and 'name' in res and 'parameters' in res:
+                functions.append(Function(name=res['name'], arguments=res['parameters']))
+        if len(functions) == 0:
+            # compat react_en
+            return super().get_toolcall(response)
+        return functions
+
+    def _format_tool_responses(
+        self,
+        assistant_content: str,
+        tool_messages,
+    ) -> Tuple[str, 'Prompt']:
+        with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content
+        if with_action:
+            return super()._format_tool_responses(assistant_content, tool_messages)
+        res = [self.eot_token]
+        for tool_message in tool_messages:
+            tool_content = tool_message['content']
+            res.append(f'{self.start_token}tool{self.end_token}\n\n{tool_content}{self.eot_token}')
+        res.append(f'{self.start_token}assistant{self.end_token}\n\n')
+        return assistant_content, res
+
+    def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
+        assert user_message is not None
+        user_content = user_message['content']
+        tool_descs = [json.dumps(tool, ensure_ascii=False, indent=4) for tool in tools]
+        new_user_content = """Given the following functions, please respond with a JSON for a function call with its proper arguments that best answers the given prompt.
+
+Respond in the format {"name": function name, "parameters": dictionary of argument name and its value}. Do not use variables.
+
+""" + '\n\n'.join(tool_descs) + f"""
+
+{user_content}"""  # noqa
+        user_message['content'] = new_user_content
+        return system or ''
+
+    def _format_tool_calls(self, tool_call_messages) -> str:
+        tool_calls = []
+        for message in tool_call_messages:
+            tool_call = self._parse_tool_call(message['content'])
+            tool_call['parameters'] = tool_call.pop('arguments')
+            tool_calls.append(json.dumps(tool_call, ensure_ascii=False))
+        return '\n'.join(tool_calls)
+
+
+class Llama4AgentTemplate(Llama3AgentTemplate):
+    eom_token = '<|eom|>'
+    start_token = '<|header_start|>'
+    end_token = '<|header_end|>'
+    eot_token = '<|eot|>'
+    toolcall_pattern = r'(.+?)<\|eom\|>'

swift/agent_template/mapping.py

@@ -0,0 +1,61 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from .deepseek_v3_1 import DeepSeekV31AgentTemplate
+from .deepseek_v4 import DeepSeekV4AgentTemplate
+from .extra import ReactGRPOAgentTemplate
+from .gemma4 import Gemma4AgentTemplate
+from .glm4 import (ChatGLM4AgentTemplate, GLM4_5AgentTemplate, GLM4_7AgentTemplate, GLM4AgentTemplate,
+                   GLM5_1AgentTemplate)
+from .hermes import HermesAgentTemplate, HunyuanHermesAgentTemplate
+from .hy_v3 import HyV3AgentTemplate, HyV3PreviewAgentTemplate
+from .kimi_k25 import KimiK25AgentTemplate
+from .llama import Llama3AgentTemplate, Llama4AgentTemplate
+from .minicpm5 import MiniCPM5AgentTemplate
+from .minimax_m2 import MinimaxM2AgentTemplate
+from .minimax_m3 import MinimaxM3AgentTemplate
+from .mistral import MistralAgentTemplate
+from .qwen import QwenEnAgentTemplate, QwenEnParallelAgentTemplate, QwenZhAgentTemplate, QwenZhParallelAgentTemplate
+from .qwen3_coder import Qwen3_5AgentTemplate, Qwen3CoderAgentTemplate
+from .react import ReactEnAgentTemplate, ReactZnAgentTemplate
+from .seed_oss import SeedAgentTemplate
+from .toolbench import ToolBenchAgentTemplate
+from .youtu import YoutuAgentTemplate
+
+agent_template_map = {
+    # ref: https://qwen.readthedocs.io/zh-cn/latest/framework/function_call.html#function-calling-templates
+    'react_en': ReactEnAgentTemplate,
+    'react_zh': ReactZnAgentTemplate,
+    # ref: https://github.com/QwenLM/Qwen-Agent/blob/main/qwen_agent/llm/fncall_prompts/qwen_fncall_prompt.py
+    'qwen_en': QwenEnAgentTemplate,
+    'qwen_zh': QwenZhAgentTemplate,
+    'qwen_en_parallel': QwenEnParallelAgentTemplate,
+    'qwen_zh_parallel': QwenZhParallelAgentTemplate,
+    'qwen3_coder': Qwen3CoderAgentTemplate,
+    'qwen3_5': Qwen3_5AgentTemplate,
+    'hermes': HermesAgentTemplate,
+    'hunyuan_hermes': HunyuanHermesAgentTemplate,
+    'hy_v3_preview': HyV3PreviewAgentTemplate,
+    'hy_v3': HyV3AgentTemplate,
+    'toolbench': ToolBenchAgentTemplate,  # ref: https://modelscope.cn/datasets/swift/ToolBench
+    'chatglm4': ChatGLM4AgentTemplate,
+    'glm4': GLM4AgentTemplate,  # ref: https://modelscope.cn/models/ZhipuAI/GLM-4-9B-0414
+    'glm4_5': GLM4_5AgentTemplate,
+    'glm4_7': GLM4_7AgentTemplate,
+    'glm5_1': GLM5_1AgentTemplate,
+    'llama3': Llama3AgentTemplate,
+    'llama4': Llama4AgentTemplate,
+    # ref: https://huggingface.co/deepseek-ai/DeepSeek-V3.1
+    'deepseek_v3_1': DeepSeekV31AgentTemplate,
+    # ref: https://modelscope.cn/models/deepseek-ai/DeepSeek-V4-Flash
+    'deepseek_v4': DeepSeekV4AgentTemplate,
+    'minimax_m2': MinimaxM2AgentTemplate,
+    'minimax_m3': MinimaxM3AgentTemplate,
+    'seed_oss': SeedAgentTemplate,
+    # ref: https://modelscope.cn/models/google/gemma-4-12B-it
+    'gemma4': Gemma4AgentTemplate,
+    # extra
+    'react_grpo': ReactGRPOAgentTemplate,
+    'mistral': MistralAgentTemplate,
+    'youtu': YoutuAgentTemplate,
+    'kimi_k25': KimiK25AgentTemplate,
+    'minicpm5': MiniCPM5AgentTemplate,
+}

swift/agent_template/minicpm5.py

@@ -0,0 +1,114 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import json
+import re
+from typing import List, Optional, Tuple, Union
+
+from swift.infer_engine import Function
+from swift.template import Prompt
+from .base import BaseAgentTemplate
+
+
+class MiniCPM5AgentTemplate(BaseAgentTemplate):
+    """Agent template for MiniCPM5 models using XML-based function calling format.
+
+    Tool call format:
+        <function name="function-name"><param name="param-name">param-value</param></function>
+
+    Tool response format:
+        <tool_response>
+        response_content
+        </tool_response>
+    """
+
+    def get_toolcall(self, response: str) -> List[Function]:
+        # Match <function name="...">...</function> blocks
+        func_pattern = re.compile(r'<function\s+name="([^"]+)">(.*?)</function>', re.DOTALL)
+        param_pattern = re.compile(r'<param\s+name="([^"]+)">'
+                                   r'(?:<!\[CDATA\[(.*?)\]\]>|([^<]*))'
+                                   r'</param>', re.DOTALL)
+
+        functions = []
+        for func_match in func_pattern.finditer(response):
+            func_name = func_match.group(1)
+            func_body = func_match.group(2)
+            arguments = {}
+            for param_match in param_pattern.finditer(func_body):
+                param_name = param_match.group(1)
+                # CDATA value or plain value
+                param_value = param_match.group(2) if param_match.group(2) is not None else param_match.group(3)
+                # Try to parse as JSON value (number, bool, etc.)
+                try:
+                    param_value = json.loads(param_value)
+                except (json.JSONDecodeError, ValueError):
+                    pass
+                arguments[param_name] = param_value
+            functions.append(Function(name=func_name, arguments=arguments))
+
+        if len(functions) == 0:
+            # Fallback to ReAct-style parsing
+            return super().get_toolcall(response)
+        return functions
+
+    def _get_tool_responses(self, tool_messages):
+        res_tool = []
+        for tool_message in tool_messages:
+            tool_content = tool_message['content']
+            res_tool.append(f'<tool_response>\n{tool_content}\n</tool_response>')
+        return '\n'.join(res_tool)
+
+    def _format_tool_responses(
+        self,
+        assistant_content: str,
+        tool_messages,
+    ) -> Tuple[str, 'Prompt']:
+        with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content
+        if with_action:
+            return super()._format_tool_responses(assistant_content, tool_messages)
+        if hasattr(self, 'template_meta'):
+            prompt = self.template_meta.prompt
+            chat_sep = self.template_meta.chat_sep
+        else:
+            prompt = ['<|im_start|>user\n{{QUERY}}<|im_end|>\n<|im_start|>assistant\n']
+            chat_sep = ['<|im_end|>\n']
+        res = chat_sep.copy()
+        total_tool = self._get_tool_responses(tool_messages)
+        for context in prompt:
+            if isinstance(context, str):
+                context = context.replace('{{QUERY}}', total_tool)
+            res.append(context)
+        return assistant_content, res
+
+    def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
+        tool_descs = [json.dumps(self.wrap_tool(tool), ensure_ascii=False) for tool in tools]
+        system = system or ''
+        if system:
+            system = f'{system}\n\n'
+        return (f'{system}# Tools\n\n'
+                'You are provided with function signatures within <tools></tools> XML tags:\n'
+                '<tools>\n' + '\n'.join(tool_descs) + '\n</tools>\n\n'
+                'Tool usage guidelines:\n'
+                '- You may call zero or more functions. If no function calls are needed, '
+                'just answer normally and do not include any <function ... </function>.\n'
+                '- When calling a function, return an XML object within <function ... </function> using:\n'
+                '<function name="function-name"><param name="param-name">param-value</param></function>\n'
+                '- param-value may be multi-line. If it contains <, & or newline characters, '
+                'wrap it in a CDATA block: <param name="param-name"><![CDATA[...multi-line value...]]></param>')
+
+    def _format_tool_calls(self, tool_call_messages) -> str:
+        tool_calls = []
+        for message in tool_call_messages:
+            tool_call = self._parse_tool_call(message['content'])
+            name = tool_call['name']
+            arguments = tool_call['arguments']
+            params_xml = ''
+            if isinstance(arguments, dict):
+                for param_name, param_value in arguments.items():
+                    value_str = param_value if isinstance(param_value, str) else json.dumps(
+                        param_value, ensure_ascii=False)
+                    if isinstance(param_value, str) and ('<' in param_value or '&' in param_value
+                                                         or '\n' in param_value):
+                        params_xml += f'<param name="{param_name}"><![CDATA[{value_str}]]></param>'
+                    else:
+                        params_xml += f'<param name="{param_name}">{value_str}</param>'
+            tool_calls.append(f'<function name="{name}">{params_xml}</function>')
+        return '\n'.join(tool_calls)

swift/agent_template/minimax_m2.py

@@ -0,0 +1,171 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import json
+import re
+from typing import List, Optional, Tuple, Union
+
+from swift.infer_engine import Function
+from swift.template import Prompt
+from .base import BaseAgentTemplate
+
+
+class MinimaxM2AgentTemplate(BaseAgentTemplate):
+    """
+    Agent template for MiniMax-M2 series models.
+
+    This template handles tool calling in MiniMax's XML-based format:
+    <minimax:tool_call>
+    <invoke name="tool-name">
+    <parameter name="param-key">param-value</parameter>
+    </invoke>
+    </minimax:tool_call>
+    """
+
+    def get_toolcall(self, response: str) -> List[Function]:
+        """
+        Extract tool calls from MiniMax response format.
+
+        Format:
+        <minimax:tool_call>
+        <invoke name="tool-name">
+        <parameter name="param-key">param-value</parameter>
+        </invoke>
+        </minimax:tool_call>
+        """
+        functions = []
+
+        # Find all tool_call blocks
+        tool_call_blocks = re.findall(r'<minimax:tool_call>(.*?)</minimax:tool_call>', response, re.DOTALL)
+
+        for block in tool_call_blocks:
+            # Find all invoke blocks within the tool_call
+            invoke_blocks = re.findall(r'<invoke name="([^"]+)">(.*?)</invoke>', block, re.DOTALL)
+
+            for tool_name, params_block in invoke_blocks:
+                # Extract parameters
+                params = {}
+                param_matches = re.findall(r'<parameter name="([^"]+)">(.*?)</parameter>', params_block, re.DOTALL)
+
+                for param_name, param_value in param_matches:
+                    param_value = param_value.strip()
+                    # Try to parse as JSON if it looks like a JSON structure
+                    parsed_value = self._parse_json(param_value)
+                    params[param_name] = parsed_value if parsed_value is not None else param_value
+
+                functions.append(Function(name=tool_name, arguments=params))
+
+        # Fallback to react format if no functions found
+        if len(functions) == 0:
+            return super().get_toolcall(response)
+
+        return functions
+
+    def _format_tool_responses(
+        self,
+        assistant_content: str,
+        tool_messages,
+    ) -> Tuple[str, 'Prompt']:
+        """
+        Format tool execution results in MiniMax format.
+
+        Tool responses are wrapped in <response></response> tags.
+        """
+        # Check if using react format
+        with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content
+        if with_action:
+            return super()._format_tool_responses(assistant_content, tool_messages)
+
+        # Use template meta if available
+        if hasattr(self, 'template_meta'):
+            prompt = self.template_meta.prompt.copy()
+            chat_sep = self.template_meta.chat_sep
+            for i in range(len(prompt)):
+                if isinstance(prompt[i], str):
+                    prompt[i] = prompt[i].replace('user', 'tool')
+        else:
+            # Default format based on the Jinja2 template
+            prompt = [']~b]tool\n{{QUERY}}[e~[\n']
+            chat_sep = ['[e~[\n']
+
+        res = chat_sep.copy()
+
+        # Format tool responses
+        tool_responses = []
+        for tool_message in tool_messages:
+            tool_content = tool_message['content']
+            tool_responses.append(f'<response>{tool_content}</response>')
+
+        total_tool = '\n'.join(tool_responses)
+
+        for context in prompt:
+            if isinstance(context, str):
+                context = context.replace('{{QUERY}}', total_tool)
+            res.append(context)
+
+        return assistant_content, res
+
+    def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
+        """
+        Format tools in MiniMax format with JSONSchema and XML invocation examples.
+        """
+        # Parse tools to JSONSchema format
+        tool_schemas = []
+        for tool in tools:
+            tool = self.unwrap_tool(tool)
+            tool_schemas.append(json.dumps(tool, ensure_ascii=False))
+
+        system = system or ''
+
+        return f"""{system}
+
+# Tools
+You may call one or more tools to assist with the user query.
+Here are the tools available in JSONSchema format:
+
+<tools>
+""" + '\n'.join(f'<tool>{schema}</tool>' for schema in tool_schemas) + """
+</tools>
+
+When making tool calls, use XML format to invoke tools and pass parameters:
+
+<minimax:tool_call>
+<invoke name="tool-name-1">
+<parameter name="param-key-1">param-value-1</parameter>
+<parameter name="param-key-2">param-value-2</parameter>
+...
+</invoke>
+</minimax:tool_call>"""
+
+    def _format_tool_calls(self, tool_call_messages):
+        """
+        Format tool call messages into MiniMax XML format.
+
+        Args:
+            tool_call_messages: List of messages containing tool call information.
+
+        Returns:
+            Formatted string with tool calls in MiniMax XML format.
+        """
+        tool_calls = []
+
+        for message in tool_call_messages:
+            tool_call = self._parse_tool_call(message['content'])
+            name = tool_call['name']
+            arguments = tool_call['arguments']
+
+            # Build parameter list
+            params = []
+            for key, value in arguments.items():
+                # Convert value to JSON string if it's not a string
+                if not isinstance(value, str):
+                    value = json.dumps(value, ensure_ascii=False)
+                params.append(f'<parameter name="{key}">{value}</parameter>')
+
+            # Build invoke block
+            invoke_block = f'<invoke name="{name}">\n' + '\n'.join(params) + '\n</invoke>'
+            tool_calls.append(invoke_block)
+
+        # Wrap all invocations in tool_call tags
+        if tool_calls:
+            return '<minimax:tool_call>\n' + '\n'.join(tool_calls) + '\n</minimax:tool_call>'
+
+        return ''

swift/agent_template/minimax_m3.py

@@ -0,0 +1,235 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import json
+import re
+from typing import Any, List, Optional, Tuple, Union
+
+from swift.infer_engine import Function
+from swift.template import Prompt
+from .base import BaseAgentTemplate
+
+# Special token used as a namespace prefix for every XML tag in MiniMax-M3
+# tool_call payloads.
+NS_TOKEN = ']<]minimax[>['
+TOOLCALL_BEGIN_TOKEN = NS_TOKEN + '<tool_call>'
+TOOLCALL_END_TOKEN = NS_TOKEN + '</tool_call>'
+
+
+def _to_xml(val: Any, ns: str = NS_TOKEN) -> str:
+    """Recursive XML renderer mirroring the ``to_xml`` macro in MiniMax-M3's
+    ``chat_template.jinja``.
+
+    ``None`` values are intentionally omitted (consistent with the upstream
+    convention that drops ``None`` parameters rather than emitting a literal
+    ``null`` string).
+    """
+    if val is None:
+        return ''
+    if isinstance(val, dict):
+        parts = []
+        for k, v in val.items():
+            if v is None:
+                continue
+            parts.append(f'{ns}<{k}>{_to_xml(v, ns)}{ns}</{k}>')
+        return ''.join(parts)
+    if isinstance(val, (list, tuple)):
+        parts = []
+        for item in val:
+            parts.append(f'{ns}<item>{_to_xml(item, ns)}{ns}</item>')
+        return ''.join(parts)
+    if isinstance(val, bool):
+        return json.dumps(val)
+    return str(val)
+
+
+_NS = re.escape(NS_TOKEN)
+_TC_BEGIN = re.escape(TOOLCALL_BEGIN_TOKEN)
+_TC_END = re.escape(TOOLCALL_END_TOKEN)
+# Match any opening tag like ]<]minimax[>[<key> (excluding closing/invoke/tool_call)
+_INVOKE_RE = re.compile(rf'{_NS}<invoke\s+name="([^"]+)">(.*?){_NS}</invoke>', re.DOTALL)
+_TOOLCALL_RE = re.compile(rf'{_TC_BEGIN}(.*?){_TC_END}', re.DOTALL)
+
+
+def _parse_xml_value(content: str) -> Any:
+    """Parse an XML fragment produced by ``to_xml`` back into a Python value.
+
+    The expected fragments use ``NS_TOKEN`` as a tag prefix. The function
+    handles nested ``<item>`` lists, dict-like ``<key>...</key>`` structures
+    and falls back to a stripped string for primitive payloads.
+    """
+    content = content.strip()
+    if not content:
+        return ''
+
+    # Try list of items first (heuristic: starts with `<item>`).
+    if content.startswith(f'{NS_TOKEN}<item>'):
+        items = []
+        for inner in _iter_tagged(content, 'item'):
+            items.append(_parse_xml_value(inner))
+        return items
+
+    # Try mapping (heuristic: starts with a NS_TOKEN<tag>).
+    if content.startswith(NS_TOKEN + '<'):
+        result: dict = {}
+        for key, inner in _iter_keyed(content):
+            result[key] = _parse_xml_value(inner)
+        if result:
+            return result
+
+    # Primitive fallback. Try JSON (booleans / numbers) before raw text.
+    try:
+        return json.loads(content)
+    except Exception:
+        return content
+
+
+def _iter_tagged(content: str, tag: str):
+    pattern = re.compile(rf'{_NS}<{re.escape(tag)}>(.*?){_NS}</{re.escape(tag)}>', re.DOTALL)
+    for m in pattern.finditer(content):
+        yield m.group(1)
+
+
+def _iter_keyed(content: str):
+    """Iterate ``(tag_name, inner_content)`` for top-level NS-prefixed tags."""
+    cursor = 0
+    n = len(content)
+    open_pat = re.compile(rf'{_NS}<([^/!?\s>]+)>')
+    while cursor < n:
+        m = open_pat.search(content, cursor)
+        if not m:
+            return
+        name = m.group(1)
+        end_marker = f'{NS_TOKEN}</{name}>'
+        # Match nested same-name tags by counting depth.
+        depth = 1
+        scan = m.end()
+        open_marker = f'{NS_TOKEN}<{name}>'
+        while depth > 0 and scan < n:
+            next_open = content.find(open_marker, scan)
+            next_close = content.find(end_marker, scan)
+            if next_close == -1:
+                return
+            if next_open != -1 and next_open < next_close:
+                depth += 1
+                scan = next_open + len(open_marker)
+            else:
+                depth -= 1
+                scan = next_close + len(end_marker)
+        inner = content[m.end():scan - len(end_marker)]
+        yield name, inner
+        cursor = scan
+
+
+class MinimaxM3AgentTemplate(BaseAgentTemplate):
+    """Agent template for MiniMax-M3 series multimodal models.
+
+    Tool calls follow this XML-with-namespace format:
+
+        ]<]minimax[>[<tool_call>
+        ]<]minimax[>[<invoke name="tool-name">
+        ]<]minimax[>[<param-1>value-1]<]minimax[>[</param-1>
+        ]<]minimax[>[<param-2>]<]minimax[>[<item>...]<]minimax[>[</item>]<]minimax[>[</param-2>
+        ]<]minimax[>[</invoke>
+        ]<]minimax[>[</tool_call>
+
+    Tool responses are wrapped in ``<response>...</response>`` inside a
+    ``]~b]tool`` slot.
+    """
+
+    def get_toolcall(self, response: str) -> List[Function]:
+        functions: List[Function] = []
+        for tc_block in _TOOLCALL_RE.findall(response):
+            for tool_name, params_block in _INVOKE_RE.findall(tc_block):
+                arguments = {}
+                for key, inner in _iter_keyed(params_block):
+                    arguments[key] = _parse_xml_value(inner)
+                functions.append(Function(name=tool_name, arguments=arguments))
+
+        if not functions:
+            return super().get_toolcall(response)
+        return functions
+
+    def _format_tool_responses(
+        self,
+        assistant_content: str,
+        tool_messages,
+    ) -> Tuple[str, 'Prompt']:
+        with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content
+        if with_action:
+            return super()._format_tool_responses(assistant_content, tool_messages)
+
+        if hasattr(self, 'template_meta'):
+            prompt = self.template_meta.prompt.copy()
+            chat_sep = self.template_meta.chat_sep
+            for i in range(len(prompt)):
+                if isinstance(prompt[i], str):
+                    prompt[i] = prompt[i].replace('user', 'tool')
+        else:
+            prompt = [']~b]tool\n{{QUERY}}[e~[\n]~b]ai\n']
+            chat_sep = ['[e~[\n']
+
+        res = chat_sep.copy() if chat_sep else []
+        tool_responses = []
+        for tool_message in tool_messages:
+            tool_content = tool_message['content']
+            tool_responses.append(f'<response>{tool_content}</response>')
+        total_tool = '\n'.join(tool_responses)
+
+        for context in prompt:
+            if isinstance(context, str):
+                context = context.replace('{{QUERY}}', total_tool)
+            res.append(context)
+        return assistant_content, res
+
+    def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
+        tool_schemas = []
+        for tool in tools:
+            tool = self.unwrap_tool(tool)
+            tool_schemas.append(json.dumps(tool, ensure_ascii=False))
+
+        system = system or ''
+        tools_xml = '\n'.join(f'<tool>{schema}</tool>' for schema in tool_schemas)
+        # Mirror the example block produced by chat_template.jinja so the
+        # in-context format hint matches inference time exactly.
+        # Note: jinja emits 'Example:\n' then '\n' before the tool_call_begin
+        # token, which renders as two consecutive newlines.
+        example = (f'\n\n{TOOLCALL_BEGIN_TOKEN}\n'
+                   f'{NS_TOKEN}<invoke name="tool-name-1">'
+                   f'{NS_TOKEN}<param-1>value-1{NS_TOKEN}</param-1>'
+                   f'{NS_TOKEN}<param-2>'
+                   f'{NS_TOKEN}<item>'
+                   f'{NS_TOKEN}<key-a>val-a{NS_TOKEN}</key-a>'
+                   f'{NS_TOKEN}<key-b>val-b{NS_TOKEN}</key-b>'
+                   f'{NS_TOKEN}</item>'
+                   f'{NS_TOKEN}</param-2>'
+                   f'{NS_TOKEN}</invoke>\n'
+                   f'{NS_TOKEN}<invoke name="tool-name-2">'
+                   f'{NS_TOKEN}<param-1>value-1{NS_TOKEN}</param-1>'
+                   f'{NS_TOKEN}</invoke>\n'
+                   f'{TOOLCALL_END_TOKEN}')
+
+        return (f'{system}\n\n# Tools\n'
+                'You may call one or more tools to assist with the user query.\n'
+                'Here are the tools available in JSONSchema format:\n'
+                f'\n<tools>\n{tools_xml}\n</tools>\n\n'
+                f'To call tools, wrap all invocations in a single {TOOLCALL_BEGIN_TOKEN}{TOOLCALL_END_TOKEN} '
+                'block. Parameter values containing nested objects or arrays are recursively expanded into '
+                f'XML elements. Example:{example}')
+
+    def _format_tool_calls(self, tool_call_messages) -> str:
+        invocations = []
+        for message in tool_call_messages:
+            tool_call = self._parse_tool_call(message['content'])
+            name = tool_call['name']
+            arguments = tool_call['arguments'] or {}
+
+            param_parts = [f'{NS_TOKEN}<invoke name="{name}">']
+            for k, v in arguments.items():
+                if v is None:
+                    continue
+                param_parts.append(f'{NS_TOKEN}<{k}>{_to_xml(v, NS_TOKEN)}{NS_TOKEN}</{k}>')
+            param_parts.append(f'{NS_TOKEN}</invoke>')
+            invocations.append(''.join(param_parts))
+
+        if not invocations:
+            return ''
+        return f'{TOOLCALL_BEGIN_TOKEN}\n' + '\n'.join(invocations) + f'\n{TOOLCALL_END_TOKEN}'

swift/agent_template/mistral.py

@@ -0,0 +1,67 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import json
+import re
+from typing import List, Optional, Tuple, Union
+
+from swift.infer_engine import Function
+from swift.template import Prompt
+from .base import BaseAgentTemplate
+
+
+class MistralAgentTemplate(BaseAgentTemplate):
+
+    def get_toolcall(self, response: str) -> List[Function]:
+        res_list = re.findall(r'\[TOOL_CALLS\]\[(.*?)\]</s>', response, re.DOTALL)
+        if not res_list:
+            return []
+        res_list = res_list[0].strip().split('\n')
+        functions = []
+        for res_str in res_list:
+            parsed_res = self._parse_json(res_str)
+            if isinstance(parsed_res, dict):
+                parsed_res = [parsed_res]  # Handle single tool call
+            if isinstance(parsed_res, list):
+                for tool_call in parsed_res:
+                    if isinstance(tool_call, dict) and 'name' in tool_call and 'arguments' in tool_call:
+                        functions.append(Function(name=tool_call['name'], arguments=tool_call['arguments']))
+        if len(functions) == 0:
+            # compat react_en
+            return super().get_toolcall(response)
+        return functions
+
+    def _format_tool_responses(
+        self,
+        assistant_content: str,
+        tool_messages,
+    ) -> Tuple[str, 'Prompt']:
+        if not hasattr(self, 'template_meta'):
+            raise ValueError('MistralAgentTemplate requires template_meta to be registered')
+        prompt = self.template_meta.prompt
+        chat_sep = self.template_meta.chat_sep
+
+        res = chat_sep.copy()
+        res_tool = []
+        for tool_message in tool_messages:
+            tool_content = tool_message['content']
+            # append `[TOOL_RESULTS]{"content": {{ .Content }}}[/TOOL_RESULTS]` to res_tool
+            res_tool.append(f'[TOOL_RESULTS]{json.dumps({"content": tool_content}, ensure_ascii=False)}[/TOOL_RESULTS]')
+        total_tool = '\n'.join(res_tool)
+        for context in prompt:
+            if isinstance(context, str):
+                context = context.replace('{{QUERY}}', total_tool)
+            res.append(context)
+        return assistant_content, res
+
+    def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
+        tool_descs = [json.dumps(self.wrap_tool(tool), ensure_ascii=False) for tool in tools]
+        system = system or ''
+        return f"""{system}[AVAILABLE_TOOLS]{' '.join(tool_descs)}[/AVAILABLE_TOOLS]"""
+
+    def _format_tool_calls(self, tool_call_messages):
+        tool_calls = []
+        for message in tool_call_messages:
+            # needs `{'name': name, 'arguments': arguments}`, which self._parse_tool_call
+            # satisfies
+            tool_call = self._parse_tool_call(message['content'])
+            tool_calls.append(json.dumps(tool_call, ensure_ascii=False))
+        return f'[TOOL_CALLS][\n{chr(10).join(tool_calls)}\n]</s>'  # check if need `</s>` at end

swift/agent_template/qwen.py

@@ -0,0 +1,134 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from typing import List, Optional, Union
+
+from .base import AgentKeyword, BaseAgentTemplate
+
+keyword = AgentKeyword(
+    action='✿FUNCTION✿:',
+    action_input='✿ARGS✿:',
+    observation='✿RESULT✿:',
+)
+
+
+class QwenEnAgentTemplate(BaseAgentTemplate):
+    keyword = keyword
+
+    def _get_tool_names_descs(self, tools):
+        tool_names = []
+        tool_descs = []
+        for tool in tools:
+            tool_desc = self._parse_tool(tool, 'en')
+            tool_names.append(tool_desc.name_for_model)
+            tool_descs.append(f'### {tool_desc.name_for_human}\n\n'
+                              f'{tool_desc.name_for_model}: {tool_desc.description_for_model} '
+                              f'Parameters: {tool_desc.parameters} {tool_desc.args_format}')
+        return tool_names, tool_descs
+
+    def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
+        tool_names, tool_descs = self._get_tool_names_descs(tools)
+        system = system or ''
+        return f"""{system}
+
+# Tools
+
+## You have access to the following tools:
+
+""" + '\n\n'.join(tool_descs) + f"""
+
+## When you need to call a tool, please insert the following command in your reply, which can be called zero or multiple times according to your needs:
+
+✿FUNCTION✿: The tool to use, should be one of [{','.join(tool_names)}]
+✿ARGS✿: The input of the tool
+✿RESULT✿: Tool results
+✿RETURN✿: Reply based on tool results. Images need to be rendered as ![](url)"""  # noqa
+
+
+class QwenZhAgentTemplate(BaseAgentTemplate):
+    keyword = keyword
+
+    def _get_tool_names_descs(self, tools):
+        tool_names = []
+        tool_descs = []
+        for tool in tools:
+            tool_desc = self._parse_tool(tool, 'zh')
+            tool_names.append(tool_desc.name_for_model)
+            tool_descs.append(f'### {tool_desc.name_for_human}\n\n'
+                              f'{tool_desc.name_for_model}: {tool_desc.description_for_model} '
+                              f'输入参数：{tool_desc.parameters} {tool_desc.args_format}')
+        return tool_names, tool_descs
+
+    def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
+        tool_names, tool_descs = self._get_tool_names_descs(tools)
+        system = system or ''
+        return f"""{system}
+
+# 工具
+
+## 你拥有如下工具：
+
+""" + '\n\n'.join(tool_descs) + f"""
+
+## 你可以在回复中插入零次、一次或多次以下命令以调用工具：
+
+✿FUNCTION✿: 工具名称，必须是[{','.join(tool_names)}]之一。
+✿ARGS✿: 工具输入
+✿RESULT✿: 工具结果
+✿RETURN✿: 根据工具结果进行回复，需将图片用![](url)渲染出来"""  # noqa
+
+
+class QwenEnParallelAgentTemplate(QwenEnAgentTemplate):
+
+    def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
+        tool_names, tool_descs = self._get_tool_names_descs(tools)
+        system = system or ''
+        return f"""{system}
+
+# Tools
+
+## You have access to the following tools:
+
+""" + '\n\n'.join(tool_descs) + f"""
+
+## Insert the following command in your reply when you need to call N tools in parallel:
+
+✿FUNCTION✿: The name of tool 1, should be one of [{','.join(tool_names)}]
+✿ARGS✿: The input of tool 1
+✿FUNCTION✿: The name of tool 2
+✿ARGS✿: The input of tool 2
+...
+✿FUNCTION✿: The name of tool N
+✿ARGS✿: The input of tool N
+✿RESULT✿: The result of tool 1
+✿RESULT✿: The result of tool 2
+...
+✿RESULT✿: he result of tool N
+✿RETURN✿: Reply based on tool results. Images need to be rendered as ![](url)"""  # noqa
+
+
+class QwenZhParallelAgentTemplate(QwenZhAgentTemplate):
+
+    def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
+        tool_names, tool_descs = self._get_tool_names_descs(tools)
+        system = system or ''
+        return f"""{system}
+
+# 工具
+
+## 你拥有如下工具：
+
+""" + '\n\n'.join(tool_descs) + f"""
+
+## 你可以在回复中插入以下命令以并行调用N个工具：
+
+✿FUNCTION✿: 工具1的名称，必须是[{','.join(tool_names)}]之一
+✿ARGS✿: 工具1的输入
+✿FUNCTION✿: 工具2的名称
+✿ARGS✿: 工具2的输入
+...
+✿FUNCTION✿: 工具N的名称
+✿ARGS✿: 工具N的输入
+✿RESULT✿: 工具1的结果
+✿RESULT✿: 工具2的结果
+...
+✿RESULT✿: 工具N的结果
+✿RETURN✿: 根据工具结果进行回复，需将图片用![](url)渲染出来"""  # noqa

swift/agent_template/qwen3_coder.py

@@ -0,0 +1,200 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import json
+import re
+from typing import List, Optional, Union
+
+from swift.infer_engine import Function
+from .hermes import HermesAgentTemplate
+
+
+def render_extra_keys(obj, handled_keys):
+    """Helper function to render extra keys not explicitly handled"""
+    result = ''
+    if isinstance(obj, dict):
+        for key, value in obj.items():
+            if key not in handled_keys:
+                result += f'\n<{key}>{json.dumps(value, ensure_ascii=False)}</{key}>'
+    return result
+
+
+TOOL_DESC_SUFFIX = (
+    '</tools>\n\nIf you choose to call a function ONLY reply in the following format with '
+    'NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\n'
+    'value_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\n'
+    'that can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\n'
+    'Reminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> '
+    'block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n'
+    '- You may provide optional reasoning for your function call in natural language BEFORE the function call, '
+    'but NOT after\n- If there is no function call available, '
+    'answer the question like normal with your current '
+    'knowledge and do not tell the user about function calls\n</IMPORTANT>')
+
+
+class Qwen3CoderAgentTemplate(HermesAgentTemplate):
+
+    @staticmethod
+    def _find_function_call(single_content: str) -> Optional[Function]:
+        single_content = single_content.strip()
+        # Check whether the complete function tag is included
+        if not single_content.startswith('<function=') or not single_content.endswith('</function>'):
+            return None
+
+        # Extract function name
+        func_name_match = re.search(r'<function=([^>]+)>', single_content)
+        if not func_name_match:
+            return None
+
+        func_name = func_name_match.group(1).strip()
+        parameters = {}
+
+        # Use regular expressions to match parameters
+        # Match any content of <parameter=name>content</parameter>
+        param_pattern = r'<parameter=([^>]+)>\s*(.*?)\s*</parameter>'
+        param_matches = re.findall(param_pattern, single_content, re.DOTALL)
+
+        for param_name, param_value in param_matches:
+            # Clear the parameter values and remove any possible additional whitespace
+            clean_value = param_value.strip()
+            parameters[param_name.strip()] = clean_value
+
+        return Function(name=func_name, arguments=json.dumps(parameters, ensure_ascii=False))
+
+    def get_toolcall(self, response: str) -> List[Function]:
+        # Extract the tool call parameters from the model's response
+        toolcall_list = re.findall(r'<tool_call>(.*?)</tool_call>', response, re.DOTALL)
+        functions = []
+        for toolcall in toolcall_list:
+            function = self._find_function_call(toolcall)
+            if function:
+                functions.append(function)
+        if len(functions) == 0:
+            # Compat react_en
+            return super(HermesAgentTemplate, self).get_toolcall(response)
+        return functions
+
+    def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
+        if system is None:
+            system = 'You are Qwen, a helpful AI assistant that can interact with a computer to solve tasks.'
+        tool_descs = [f'{system}\n\n# Tools\n\nYou have access to the following functions:\n\n<tools>']
+        for tool in tools:
+            tool_desc = ''
+
+            # Check function key
+            if isinstance(tool, dict) and 'function' in tool:
+                tool = tool['function']
+
+            # Add function name
+            tool_desc += f"<function>\n<name>{tool['name']}</name>"
+
+            # Add description if available
+            if 'description' in tool:
+                tool_desc += f"\n<description>{tool['description'].strip()}</description>"
+
+            # Add parameters section
+            tool_desc += '\n<parameters>'
+
+            # Process parameters if they exist in the expected structure
+            if ('parameters' in tool and isinstance(tool['parameters'], dict) and 'properties' in tool['parameters']
+                    and isinstance(tool['parameters']['properties'], dict)):
+
+                for param_name, param_fields in tool['parameters']['properties'].items():
+                    tool_desc += '\n<parameter>'
+                    tool_desc += f'\n<name>{param_name}</name>'
+
+                    if 'type' in param_fields:
+                        tool_desc += f"\n<type>{str(param_fields['type'])}</type>"
+
+                    if 'description' in param_fields:
+                        tool_desc += f"\n<description>{param_fields['description'].strip()}</description>"
+
+                    # Add any extra parameter fields
+                    handled_param_keys = ['name', 'type', 'description']
+                    tool_desc += render_extra_keys(param_fields, handled_param_keys)
+
+                    tool_desc += '\n</parameter>'
+            # Add any extra parameter section fields
+            handled_keys = ['type', 'properties']
+            if 'parameters' in tool:
+                tool_desc += render_extra_keys(tool['parameters'], handled_keys)
+
+            tool_desc += '\n</parameters>'
+
+            # Add any extra function fields
+            handled_keys = ['type', 'name', 'description', 'parameters']
+            tool_desc += render_extra_keys(tool, handled_keys)
+
+            tool_desc += '\n</function>'
+
+            tool_descs.append(tool_desc)
+
+        tool_descs.append(TOOL_DESC_SUFFIX)
+        tool_descs = '\n'.join(tool_descs)
+        return tool_descs
+
+    def _format_tool_calls(self, tool_call_messages):
+        result_parts = []
+        for idx, message in enumerate(tool_call_messages):
+            tool_call = self._parse_tool_call(message['content'])
+            result_parts.append(f"<tool_call>\n<function={tool_call['name']}>\n")
+            # Processing parameters (if present)
+            if 'arguments' in tool_call and tool_call['arguments']:
+                for args_name, args_value in tool_call['arguments'].items():
+                    result_parts.append(f'<parameter={args_name}>\n')
+                    # Handle different types of parameter values
+                    if isinstance(args_value, (dict, list)):
+                        # For dictionaries or lists, use json formatting
+                        args_value = json.dumps(args_value, ensure_ascii=False)
+                    else:
+                        # For other types, convert to strings
+                        args_value = str(args_value)
+                    result_parts.append(f'{args_value}\n</parameter>\n')
+            # Close tags
+            result_parts.append('</function>\n</tool_call>')
+            # ref: https://github.com/QwenLM/Qwen3-Coder/blob/0ae30f55e9d6c47ff763c334f99c135ad68915dd/qwencoder-eval/tool_calling_eval/berkeley-function-call-leaderboard/bfcl_eval/model_handler/local_inference/qwen_fc.py#L21  # noqa
+            if idx != len(tool_call_messages) - 1:
+                result_parts.append('\n')
+        return ''.join(result_parts)
+
+    def _get_tool_responses(self, tool_messages):
+        res_tool = []
+        for tool_message in tool_messages:
+            tool_content = tool_message['content']
+            res_tool.append(f'<tool_response>\n{tool_content}\n</tool_response>\n')
+        return ''.join(res_tool)
+
+
+class Qwen3_5AgentTemplate(Qwen3CoderAgentTemplate):
+
+    def _add_tool_call_prefix(self, tool_content: str, pre_message=None) -> str:
+        """Qwen3.5/3.6 jinja inserts \n\n between assistant content and <tool_call>
+        only when effective content (after think removal) is non-empty."""
+        if not pre_message or pre_message.get('role') != 'assistant':
+            return tool_content
+        content = pre_message.get('content', '')
+        if not isinstance(content, str):
+            return tool_content
+        # Mirror jinja: content.split('</think>')[-1].lstrip('\n') then content|trim
+        if '</think>' in content:
+            effective = content.split('</think>')[-1].lstrip('\n')
+        else:
+            effective = content
+        if effective.strip():
+            return '\n\n' + tool_content
+        return tool_content
+
+    def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
+        tool_descs = [json.dumps(self.wrap_tool(tool), ensure_ascii=False) for tool in tools]
+        tools_prompt = """# Tools
+
+You have access to the following functions:\n\n<tools>
+""" + '\n'.join(tool_descs) + f'\n{TOOL_DESC_SUFFIX}'
+        if system:
+            tools_prompt += f'\n\n{system}'
+        return tools_prompt
+
+    def _get_tool_responses(self, tool_messages):
+        res_tool = []
+        for tool_message in tool_messages:
+            tool_content = tool_message['content']
+            res_tool.append(f'<tool_response>\n{tool_content}\n</tool_response>')
+        return '\n'.join(res_tool)

swift/agent_template/react.py

@@ -0,0 +1,66 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from typing import List, Optional, Union
+
+from .base import BaseAgentTemplate
+
+
+class ReactEnAgentTemplate(BaseAgentTemplate):
+
+    def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
+        tool_names = []
+        tool_descs = []
+        for tool in tools:
+            tool_desc = self._parse_tool(tool, 'en')
+            tool_names.append(tool_desc.name_for_model)
+            tool_descs.append(
+                f'{tool_desc.name_for_model}: Call this tool to interact with the {tool_desc.name_for_human} API. '
+                f'What is the {tool_desc.name_for_human} API useful for? {tool_desc.description_for_model} '
+                f'Parameters: {tool_desc.parameters} {tool_desc.args_format}')
+
+        return """Answer the following questions as best you can. You have access to the following tools:
+
+""" + '\n\n'.join(tool_descs) + f"""
+
+Use the following format:
+
+Question: the input question you must answer
+Thought: you should always think about what to do
+Action: the action to take, should be one of [{','.join(tool_names)}]
+Action Input: the input to the action
+Observation: the result of the action
+... (this Thought/Action/Action Input/Observation can be repeated zero or more times)
+Thought: I now know the final answer
+Final Answer: the final answer to the original input question
+
+Begin!
+"""
+
+
+class ReactZnAgentTemplate(BaseAgentTemplate):
+
+    def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
+        tool_names = []
+        tool_descs = []
+        for tool in tools:
+            tool_desc = self._parse_tool(tool, 'zh')
+            tool_names.append(tool_desc.name_for_model)
+            tool_descs.append(f'{tool_desc.name_for_model}: 调用此工具与 {tool_desc.name_for_human} API 进行交互。'
+                              f'{tool_desc.name_for_human} 有什么用？{tool_desc.description_for_model} '
+                              f'输入参数：{tool_desc.parameters} {tool_desc.args_format}')
+        return """尽可能地回答以下问题。你可以使用以下工具:
+
+""" + '\n\n'.join(tool_descs) + f"""
+
+请按照以下格式进行:
+
+Question: 需要你回答的输入问题
+Thought: 你应该总是思考该做什么
+Action: 需要使用的工具，应该是[{','.join(tool_names)}]中的一个
+Action Input: 传入工具的内容
+Observation: 行动的结果
+... (这个Thought/Action/Action Input/Observation可以重复N次)
+Thought: 我现在知道最后的答案
+Final Answer: 对原始输入问题的最终答案
+
+现在开始！
+"""

swift/agent_template/seed_oss.py

@@ -0,0 +1,153 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import json
+import re
+from typing import List, Optional, Tuple, Union
+
+from swift.infer_engine import Function
+from swift.template import Prompt
+from .base import BaseAgentTemplate
+
+
+class SeedAgentTemplate(BaseAgentTemplate):
+    TOOL_CALL_START = '<seed:tool_call>'
+    TOOL_CALL_END = '</seed:tool_call>'
+    FUNCTION_TAG = 'function'
+    PARAMETER_TAG = 'parameter'
+
+    _PY_TYPE_MAPPING = {
+        'string': 'str',
+        'number': 'int',
+        'integer': 'int',
+        'boolean': 'bool',
+        'array': 'list',
+    }
+
+    @staticmethod
+    def _py_type(t: str) -> str:
+        return SeedAgentTemplate._PY_TYPE_MAPPING.get(t, 'Any')
+
+    def get_toolcall(self, response: str) -> List[Function]:
+        res_list = re.findall(rf'{self.TOOL_CALL_START}(.+?){self.TOOL_CALL_END}', response, re.DOTALL)
+        if not res_list:
+            return super().get_toolcall(response)
+
+        functions = []
+        for res in res_list:
+            func_name_match = re.search(rf'<{self.FUNCTION_TAG}=([^>]+)>', res)
+            if not func_name_match:
+                continue
+
+            func_name = func_name_match.group(1)
+            param_matches = re.findall(rf'<{self.PARAMETER_TAG}=([^>]+)>(.*?)</{self.PARAMETER_TAG}>', res, re.DOTALL)
+            arguments = {name: value for name, value in param_matches}
+            functions.append(Function(name=func_name, arguments=arguments))
+
+        return functions
+
+    def _get_tool_responses(self, tool_messages: List[dict]) -> str:
+        responses = [f"<seed:bos>tool\n{tool_message['content']}<seed:eos>" for tool_message in tool_messages]
+        return ''.join(responses) + '<seed:bos>assistant\n'
+
+    def _format_tool_responses(
+        self,
+        assistant_content: str,
+        tool_messages: List[dict],
+    ) -> Tuple[str, 'Prompt']:
+        with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content
+        if with_action:
+            return super()._format_tool_responses(assistant_content, tool_messages)
+
+        formatted_tool_responses = self._get_tool_responses(tool_messages)
+        return assistant_content, ['<seed:eos>', formatted_tool_responses]
+
+    def _build_tool_def_string(self, tool: dict) -> str:
+        """Helper to build a single tool definition string."""
+        func = tool.get('function', {})
+        func_name = func.get('name')
+
+        if not func_name:
+            return ''
+
+        parameters = func.get('parameters', {})
+        properties = parameters.get('properties', {})
+        params = [
+            f"{name}: {self._py_type(spec.get('type', 'any'))}" for name, spec in properties.items()
+            if isinstance(spec, dict)
+        ]
+        param_str = ','.join(params)
+
+        docstring_parts = ['    """', f'    {func.get("description", "").strip()}']
+
+        if properties:
+            docstring_parts.append('\n    Args:')
+            required_params = parameters.get('required', [])
+            for name, spec in properties.items():
+                if isinstance(spec, dict):
+                    req_tag = '[必填]' if name in required_params else '[选填]'
+                    desc = spec.get('description', '')
+                    type_str = self._py_type(spec.get('type', 'any'))
+                    docstring_parts.append(f'    - {name} ({type_str}) {req_tag}: {desc}')
+
+        returns_props = func.get('returns', {}).get('properties', {})
+        if returns_props:
+            docstring_parts.append('\n    Returns:')
+            for name, spec in returns_props.items():
+                desc = spec.get('description', '')
+                type_str = self._py_type(spec.get('type', 'any'))
+                docstring_parts.append(f'    - {name} ({type_str}): {desc}')
+
+        docstring_parts.append('\n    """')
+        docstring = '\n'.join(docstring_parts)
+
+        return f'Function:\ndef {func_name}({param_str}):\n{docstring}'
+
+    def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
+        if not tools:
+            return system or ''
+
+        tool_defs = [
+            tool_def for tool in tools if (wrapped_tool := self.wrap_tool(tool)).get('type') == 'function' and (
+                tool_def := self._build_tool_def_string(wrapped_tool)) != ''
+        ]
+        tool_defs_joined = '\n\n'.join(tool_defs)
+
+        tool_call_format_instruction = (
+            '工具调用请遵循如下格式:\n'
+            f'{self.TOOL_CALL_START}\n'
+            f'<{self.FUNCTION_TAG}=example_function_name>\n'
+            f'<{self.PARAMETER_TAG}=example_parameter_1>value_1</{self.PARAMETER_TAG}>\n'
+            f'<{self.PARAMETER_TAG}=example_parameter_2>This is the value for the second parameter\n'
+            'that can span\n'
+            f'multiple lines</{self.PARAMETER_TAG}>\n'
+            f'</{self.FUNCTION_TAG}>\n'
+            f'{self.TOOL_CALL_END}')
+
+        split_token = '<seed:eos><seed:bos>system'
+
+        if system and split_token in system:
+            parts = system.split(split_token, 1)
+            return f'{parts[0]}\n\n{tool_defs_joined}\n{tool_call_format_instruction}\n{split_token}{parts[1]}'
+        else:
+            doubao_prompt = ('You are Doubao, a helpful AI assistant. '
+                             'You may call one or more functions to assist with the user query.')
+            return (f'{doubao_prompt}\n\n{tool_defs_joined}\n{tool_call_format_instruction}\n'
+                    f'{split_token}\n{system or ""}')
+
+    def _format_tool_calls(self, tool_call_messages: List[dict]) -> str:
+        formatted_calls = []
+        for message in tool_call_messages:
+            tool_call = self._parse_tool_call(message['content'])
+            func_name = tool_call['name']
+            arguments = tool_call.get('arguments', {})
+
+            call_parts = [f'<{self.FUNCTION_TAG}={func_name}>']
+            for arg_name, arg_value in arguments.items():
+                arg_value_str = arg_value if isinstance(arg_value, str) else json.dumps(arg_value, ensure_ascii=False)
+                call_parts.append(f'<{self.PARAMETER_TAG}={arg_name}>{arg_value_str}</{self.PARAMETER_TAG}>')
+
+            call_parts.append(f'</{self.FUNCTION_TAG}>')
+            call_parts_joined = '\n'.join(call_parts)
+
+            full_call = f'{self.TOOL_CALL_START}\n{call_parts_joined}\n{self.TOOL_CALL_END}'
+            formatted_calls.append(full_call)
+        return '\n'.join(formatted_calls)

swift/agent_template/toolbench.py

@@ -0,0 +1,38 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import json
+from typing import List, Optional, Union
+
+from .base import BaseAgentTemplate
+
+
+class ToolBenchAgentTemplate(BaseAgentTemplate):
+
+    def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
+        for i, tool in enumerate(tools):
+            tools[i] = self.unwrap_tool(tool)
+        tools = json.dumps(tools, ensure_ascii=False)
+        return f"""You can use many tools(functions) to do the following task.
+First I will give you the task description, and your task start.
+At each step, you need to give your thought to analyze the status now and what to do next, \
+with a function call to actually execute your step. Your output should follow this format:
+Thought:
+Action:
+Action Input:
+
+After the call, you will get the call result, and you are now in a new state.
+Then you will analyze your status now, then decide what to do next...
+After many (Thought-call) pairs, you finally perform the task, then you can give your final answer.
+Remember:
+1.the state change is irreversible, you can't go back to one of the former state, if you want to restart the task, \
+say \"I give up and restart\".
+2.All the thought is short, at most in 5 sentence.
+3.You can do more then one try, so if your plan is to continuously try some conditions, \
+you can do one of the conditions per try.
+Let's Begin!
+Task description: You should use functions to help handle the real time user queries. Remember:
+1.ALWAYS call \"Finish\" function at the end of the task. And the final answer should contain enough information \
+to show to the user,If you can't handle the task, \
+or you find that function calls always fail(the function is not valid now), \
+use function Finish->give_up_and_restart.
+2.Do not use origin tool names, use only subfunctions' names.
+Specifically, you have access to the following APIs: {tools}"""

swift/agent_template/youtu.py

@@ -0,0 +1,66 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import json
+from typing import TYPE_CHECKING, List, Optional, Tuple, Union
+
+from .hermes import HermesAgentTemplate
+
+if TYPE_CHECKING:
+    from swift.template import Prompt
+
+
+class YoutuAgentTemplate(HermesAgentTemplate):
+    """Agent template for Youtu-LLM models.
+
+    Tool calling format:
+    - Tool call: <tool_call>{"name": "function-name", "arguments": {...}}</tool_call>
+    - Tool response: <tool_response>...</tool_response>
+    """
+
+    def _get_tool_responses(self, tool_messages):
+        res_tool = []
+        for tool_message in tool_messages:
+            tool_content = tool_message['content']
+            res_tool.append(f'<tool_response>{tool_content}</tool_response>')
+        return '\n'.join(res_tool)
+
+    def _format_tool_responses(
+        self,
+        assistant_content: str,
+        tool_messages,
+    ) -> Tuple[str, 'Prompt']:
+        with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content
+        if with_action:
+            return super()._format_tool_responses(assistant_content, tool_messages)
+        # For Youtu-LLM, tool responses are placed in user message
+        if hasattr(self, 'template_meta'):
+            prompt = self.template_meta.prompt
+            chat_sep = self.template_meta.chat_sep
+        else:
+            prompt = ['<|User|>{{QUERY}}<|Assistant|>']
+            chat_sep = ['<|end_of_text|>']
+        res = chat_sep.copy()
+        total_tool = self._get_tool_responses(tool_messages)
+        for context in prompt:
+            if isinstance(context, str):
+                context = context.replace('{{QUERY}}', total_tool)
+            res.append(context)
+        return assistant_content, res
+
+    def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
+        tool_descs = [json.dumps(self.wrap_tool(tool), ensure_ascii=False) for tool in tools]
+        system = system or ''
+        if system:
+            system = f'{system}\n\n'
+        return f"""{system}<|begin_of_tool_description|>Tool calling capabilities.
+You may call one or more functions to assist with the user query. You have the following functions available:
+""" + '\n'.join([f'```json\n{desc}\n```' for desc in tool_descs]) + """
+For tool call returns, you MUST use the following format:
+<tool_call>{"name": "function-name", "arguments": {"param1": "value1", "param2": "value2"}}</tool_call>
+<|end_of_tool_description|>"""
+
+    def _format_tool_calls(self, tool_call_messages):
+        tool_calls = []
+        for message in tool_call_messages:
+            tool_call = self._parse_tool_call(message['content'])
+            tool_calls.append(f'<tool_call>{json.dumps(tool_call, ensure_ascii=False)}</tool_call>')
+        return ''.join(tool_calls)

swift/arguments/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from .app_args import AppArguments
+from .base_args import BaseArguments, DataArguments, ModelArguments, TemplateArguments, get_supported_tuners
+from .deploy_args import DeployArguments, RolloutArguments
+from .eval_args import EvalArguments
+from .export_args import ExportArguments
+from .infer_args import InferArguments
+from .pretrain_args import PretrainArguments
+from .rlhf_args import RLHFArguments
+from .sampling_args import SamplingArguments
+from .sft_args import SftArguments
+from .tuner_args import TunerArguments
+from .webui_args import WebUIArguments

swift/arguments/app_args.py

@@ -0,0 +1,57 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from dataclasses import dataclass
+from typing import Literal, Optional
+
+from swift.model import get_matched_model_meta
+from swift.template import get_template_meta
+from swift.utils import find_free_port, get_logger
+from .deploy_args import DeployArguments
+from .webui_args import WebUIArguments
+
+logger = get_logger()
+
+
+@dataclass
+class AppArguments(WebUIArguments, DeployArguments):
+    """Arguments for configuring the Web UI inference.
+
+    This dataclass inherits from WebUIArguments and DeployArguments, combining their settings to configure the user
+    interface for model inference.
+
+    Args:
+        base_url (Optional[str]): The base URL for the model deployment API, e.g., `http://localhost:8000/v1`. If set
+            to `None`, a local deployment will be used instead. Defaults to None.
+        studio_title (Optional[str]): The title for the Web UI studio. If set to `None`, the title will default to the
+            model's name. Defaults to None.
+        is_multimodal (Optional[bool]): Whether to launch the multimodal version of the application. If `None`, the
+            app will attempt to auto-detect this setting based on the model. If auto-detection is not possible, it
+            defaults to `False`. Defaults to None.
+        lang (str): Overrides the language setting for the Web UI. Defaults to 'en'.
+        verbose (bool): Whether to log detailed request information. Defaults to False.
+        stream (bool): Whether to enable streaming output for model responses. Defaults to True.
+    """
+    base_url: Optional[str] = None
+    studio_title: Optional[str] = None
+    is_multimodal: Optional[bool] = None
+
+    lang: Literal['en', 'zh'] = 'en'
+    verbose: bool = False
+    stream: bool = True
+
+    def _init_torch_dtype(self) -> None:
+        if self.base_url:
+            self.model_meta = get_matched_model_meta(self.model)
+            self.model_info = None
+            return
+        super()._init_torch_dtype()
+
+    def __post_init__(self):
+        DeployArguments.__post_init__(self)
+        self.server_port = find_free_port(self.server_port)
+        if self.model_meta:
+            if self.system is None:
+                self.system = get_template_meta(self.model_info, self.model_meta).default_system
+            if self.is_multimodal is None:
+                self.is_multimodal = self.model_meta.is_multimodal
+        if self.is_multimodal is None:
+            self.is_multimodal = False

swift/arguments/base_args/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from .base_args import BaseArguments, get_supported_tuners
+from .data_args import DataArguments
+from .generation_args import GenerationArguments
+from .model_args import ModelArguments
+from .quant_args import QuantizeArguments
+from .template_args import TemplateArguments

swift/arguments/base_args/base_args.py

@@ -0,0 +1,367 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import json
+import os
+import peft
+import shutil
+from dataclasses import dataclass, field, fields
+from packaging import version
+from typing import Any, Dict, List, Literal, Optional, Union
+
+import swift
+from swift.dataset import load_dataset
+from swift.hub import get_hub
+from swift.model import get_ckpt_dir, get_model_processor, load_by_unsloth
+from swift.ray_utils import RayArguments
+from swift.template import Template, get_template
+from swift.tuner_plugin import tuners_map
+from swift.utils import (Processor, check_json_format, get_dist_setting, get_logger, import_external_file, is_dist,
+                         is_master, json_parse_to_dict, safe_snapshot_download, set_device, use_hf_hub)
+from .data_args import DataArguments
+from .generation_args import GenerationArguments
+from .model_args import ModelArguments
+from .quant_args import QuantizeArguments
+from .template_args import TemplateArguments
+
+logger = get_logger()
+
+
+def get_supported_tuners():
+    return {'lora', 'full', 'longlora', 'adalora', 'llamapro', 'adapter', 'vera', 'boft', 'fourierft', 'reft', 'bone'
+            } | set(tuners_map.keys())
+
+
+def _patch_peft():
+    """Patch peft functions that are incompatible with SWIFT.
+
+    1. _maybe_shard_state_dict_for_tp: TP sharding is not used by SWIFT, and causes errors
+       when torch.distributed is initialized (e.g. MoE training with target_parameters).
+    2. _maybe_shard_state_dict_for_tp internal logic accesses base_layer.weight.device which
+       fails for expert modules that don't have a `weight` attribute.
+    """
+    if version.parse(peft.__version__) >= version.parse('0.19.0'):
+        from peft.utils import save_and_load
+        save_and_load._maybe_shard_state_dict_for_tp = lambda model, state_dict, adapter_name: None
+
+
+@dataclass
+class BaseArguments(GenerationArguments, QuantizeArguments, DataArguments, TemplateArguments, ModelArguments,
+                    RayArguments):
+    """BaseArguments class is a dataclass that inherits from multiple argument classes.
+
+    This class consolidates arguments from GenerationArguments, QuantizeArguments, DataArguments,
+    TemplateArguments, ModelArguments, RayArguments.
+
+    Args:
+        tuner_backend (str): The tuner backend to use. Choices are 'peft' or 'unsloth'. Default is 'peft'.
+        tuner_type (str): The tuner type. Choices include 'lora', 'full', 'longlora', 'adalora', 'llamapro',
+            'adapter', 'vera', 'boft', 'fourierft', 'reft'. Default is 'lora'.
+        adapters (List[str]): A list of adapter IDs or paths. This is typically used for inference or deployment.
+            It can also resume training by only loading adapter weights, differing from `resume_from_checkpoint`
+            which also loads optimizer states. Default is [].
+        external_plugins (List[str]): A list of external 'plugin.py' files to be registered and imported into
+            the plugin module. Default is [].
+        seed (int): The global random seed for reproducibility. Note that this does not affect `data_seed`,
+            which controls dataset randomization. Default is 42.
+        model_kwargs (Optional[str]): Additional keyword arguments for specific models, passed as a JSON string
+            (e.g., '{"key": "value"}'). It's recommended to use the same arguments for inference as for training.
+            Default is None.
+        enable_npu_model_patch (bool): Whether to enable model-related NPU patches. Default is True.
+        load_args (bool): Whether to load `args.json` from a checkpoint when using `--resume_from_checkpoint`,
+            `--model`, or `--adapters`. Defaults to True for inference/export and False for training. Usually,
+            this does not need to be modified. Default is True.
+        load_data_args (bool): If True, will also load data-related arguments from `args.json`. This is useful
+            for running inference on the same validation split used during training. Default is False.
+        packing (bool): Whether to enable packing of datasets. Default is False.
+        packing_length (Optional[int]): Length of packing. Default is None.
+        packing_num_proc (int): Number of processes used for packing, Default is 1.
+        packing_strategy (Literal['binpack', 'sequential']): Packing algorithm. 'binpack' (default) uses
+            best-fit-decreasing bin packing (reorders samples); 'sequential' uses order-preserving greedy
+            packing (next-fit: a single open pack, flushed when the next sample doesn't fit) so the sample
+            order / pack boundaries follow a sequential sampler (use packing_num_proc=1). Default is 'binpack'.
+        lazy_tokenize (Optional[bool]): Whether to enable lazy tokenization. Default is None.
+        use_hf (bool): Whether to use Hugging Face for downloading/uploading models and datasets. If False,
+            ModelScope is used. Default is False.
+        hub_token (Optional[str]): The authentication token for ModelScope or Hugging Face Hub. Default is None.
+        ddp_timeout (int): Timeout for DDP (Distributed Data Parallel) operations, in seconds. Default is 18000000.
+        ddp_backend (Optional[str]): The backend for DDP. Choices include "nccl", "gloo", "mpi", "ccl", "hccl",
+            "cncl", "mccl". If None, it will be automatically selected. Default is None.
+        ignore_args_error (bool): Whether to ignore argument errors. This is useful for compatibility with Jupyter
+            notebooks. Default is False.
+        use_swift_lora (bool): Whether to use swift lora. This is a compatible argument. Default is False.
+    """
+    tuner_backend: Literal['peft', 'unsloth'] = 'peft'
+    tuner_type: str = field(default='lora', metadata={'help': f'tuner_type choices: {list(get_supported_tuners())}'})
+    adapters: List[str] = field(default_factory=list)
+    external_plugins: List[str] = field(default_factory=list)
+    # This parameter is kept for swift3.x compatibility. Please use `external_plugins` as a replacement.
+    custom_register_path: List[str] = field(default_factory=list)
+
+    seed: int = 42
+    model_kwargs: Optional[Union[dict, str]] = None
+    enable_npu_model_patch: bool = True
+    load_args: bool = True
+    load_data_args: bool = False
+    # dataset
+    packing: bool = False
+    packing_length: Optional[int] = None
+    packing_num_proc: int = 1
+    packing_strategy: Literal['binpack', 'sequential'] = 'binpack'
+    lazy_tokenize: Optional[bool] = None
+    # hub
+    use_hf: bool = False
+    # None: use env var `MODELSCOPE_API_TOKEN`
+    hub_token: Optional[str] = field(
+        default=None, metadata={'help': 'SDK token can be found in https://modelscope.cn/my/myaccesstoken'})
+    # dist
+    ddp_timeout: int = 18000000
+    ddp_backend: Optional[str] = None
+
+    # extra
+    ignore_args_error: bool = False  # True: notebook compatibility
+    use_swift_lora: bool = False  # True for using tuner_backend == swift, don't specify this unless you know what you are doing # noqa
+
+    def _prepare_training_args(self, training_args: Dict[str, Any]) -> None:
+        pass
+
+    def _init_lazy_tokenize(self):
+        if self.lazy_tokenize is None:
+            if self.cached_dataset or self.cached_val_dataset:
+                self.lazy_tokenize = False
+            elif (self.model_meta is not None and self.model_meta.is_multimodal and not self.streaming
+                  and not self.packing and not getattr(self, 'group_by_length', False)):
+                self.lazy_tokenize = True
+            else:
+                self.lazy_tokenize = False
+            logger.info(f'Setting args.lazy_tokenize: {self.lazy_tokenize}')
+        if self.lazy_tokenize:
+            if self.packing:
+                raise ValueError('Packing and lazy_tokenize are incompatible.')
+            if self.streaming:
+                raise ValueError('Streaming and lazy_tokenize are incompatible.')
+
+    def _import_external_plugins(self):
+        if isinstance(self.external_plugins, str):
+            self.external_plugins = [self.external_plugins]
+        # swift v3.x compatibility
+        if isinstance(self.custom_register_path, str):
+            self.custom_register_path = [self.custom_register_path]
+        if self.custom_register_path:
+            self.external_plugins += self.custom_register_path
+
+        if not self.external_plugins:
+            return
+        for external_plugin in self.external_plugins:
+            import_external_file(external_plugin)
+        logger.info(f'Successfully imported external_plugins: {self.external_plugins}.')
+
+    @staticmethod
+    def _check_is_adapter(adapter_dir: str) -> bool:
+        if (os.path.exists(os.path.join(adapter_dir, 'adapter_config.json'))
+                or os.path.exists(os.path.join(adapter_dir, 'default', 'adapter_config.json'))
+                or os.path.exists(os.path.join(adapter_dir, 'reft'))):
+            return True
+        return False
+
+    def _init_adapters(self):
+        if isinstance(self.adapters, str):
+            self.adapters = [self.adapters]
+        self.adapters = [
+            safe_snapshot_download(adapter, use_hf=self.use_hf, hub_token=self.hub_token) for adapter in self.adapters
+        ]
+
+    def __post_init__(self):
+        _patch_peft()
+        self.swift_version = swift.__version__
+        if self.use_hf or use_hf_hub():
+            self.use_hf = True
+            os.environ['USE_HF'] = '1'
+        self._init_adapters()
+        self._init_ckpt_dir()
+        self._import_external_plugins()
+        self._init_model_kwargs()
+        # The Seq2SeqTrainingArguments has a property called world_size, which cannot be assigned a value.
+        self.rank, self.local_rank, self.global_world_size, self.local_world_size = get_dist_setting()
+        logger.info(f'rank: {self.rank}, local_rank: {self.local_rank}, '
+                    f'world_size: {self.global_world_size}, local_world_size: {self.local_world_size}')
+        if self.tuner_type not in tuners_map:  # build-in tuner
+            for adapter in self.adapters:
+                assert self._check_is_adapter(adapter), (
+                    f'`{adapter}` is not an adapter, please try using `--model` to pass it.')
+        ModelArguments.__post_init__(self)
+        QuantizeArguments.__post_init__(self)
+        TemplateArguments.__post_init__(self)
+        DataArguments.__post_init__(self)
+        RayArguments.__post_init__(self)
+        self._init_stream()
+        if self.max_length is None and self.model_info is not None:
+            self.max_length = self.model_info.max_model_len
+        if self.packing and self.packing_length is None:
+            self.packing_length = self.max_length
+        self._init_lazy_tokenize()
+        self.hub = get_hub(self.use_hf)
+        if self.hub.try_login(self.hub_token):
+            logger.info('hub login successful!')
+
+    def _init_model_kwargs(self):
+        """Prepare model kwargs and set them to the env"""
+        self.model_kwargs: Dict[str, Any] = json_parse_to_dict(self.model_kwargs)
+        for k, v in self.model_kwargs.items():
+            k = k.upper()
+            os.environ[k] = str(v)
+
+    @property
+    def is_adapter(self) -> bool:
+        return self.tuner_type not in {'full'}
+
+    @property
+    def supported_tuners(self):
+        return get_supported_tuners()
+
+    @property
+    def adapters_can_be_merged(self):
+        return {'lora', 'longlora', 'llamapro', 'adalora'}
+
+    @classmethod
+    def from_pretrained(cls, checkpoint_dir: str):
+        self = super().__new__(cls)
+        self.load_data_args = True
+        self.ckpt_dir = checkpoint_dir
+        self.load_args_from_ckpt()
+        all_keys = list(f.name for f in fields(BaseArguments))
+        for key in all_keys:
+            if not hasattr(self, key):
+                setattr(self, key, None)
+        return self
+
+    def _init_ckpt_dir(self, adapters=None):
+        # compat megatron
+        model = self.model or getattr(self, 'mcore_model', None)
+        adapters = adapters or self.adapters or getattr(self, 'mcore_adapter', None)
+        if isinstance(adapters, str):
+            adapters = [adapters]
+        self.ckpt_dir = get_ckpt_dir(model, adapters)
+        if self.ckpt_dir and self.load_args:
+            self.load_args_from_ckpt()
+
+    def load_args_from_ckpt(self) -> None:
+        args_path = os.path.join(self.ckpt_dir, 'args.json')
+        assert os.path.exists(args_path), f'args_path: {args_path}'
+        with open(args_path, 'r', encoding='utf-8') as f:
+            old_args = json.load(f)
+        force_load_keys = [
+            # base_args
+            'tuner_type',
+            # model_args
+            'task_type',
+            # quant_args
+            'bnb_4bit_quant_type',
+            'bnb_4bit_use_double_quant',
+        ]
+        # If the current value is None or an empty list and it is among the following keys
+        load_keys = [
+            'external_plugins',
+            # model_args
+            'model',
+            'model_type',
+            'model_revision',
+            'torch_dtype',
+            'attn_impl',
+            'experts_impl',
+            'new_special_tokens',
+            'num_labels',
+            'problem_type',
+            'rope_scaling',
+            'max_model_len',
+            # quant_args
+            'quant_method',
+            'quant_bits',
+            'hqq_axis',
+            'bnb_4bit_compute_dtype',
+            # template_args
+            'template',
+            'system',
+            'truncation_strategy',
+            'agent_template',
+            'norm_bbox',
+            'use_chat_template',
+            'response_prefix',
+        ]
+        data_keys = list(f.name for f in fields(DataArguments))
+        swift_version = old_args.get('swift_version')
+        if swift_version is None or version.parse(swift_version) < version.parse('4.0.0.dev'):
+            load_keys.remove('model_type')
+        for key, old_value in old_args.items():
+            if old_value is None:
+                continue
+            if key in force_load_keys or self.load_data_args and key in data_keys:
+                setattr(self, key, old_value)
+            value = getattr(self, key, None)
+            if key in load_keys and (value is None or isinstance(value, (list, tuple)) and len(value) == 0):
+                setattr(self, key, old_value)
+        logger.info(f'Successfully loaded {args_path}.')
+
+    def save_args(self, output_dir=None) -> None:
+        if is_master():
+            output_dir = output_dir or self.output_dir
+            os.makedirs(output_dir, exist_ok=True)
+            fpath = os.path.join(output_dir, 'args.json')
+            logger.info(f'The {self.__class__.__name__} will be saved in: {fpath}')
+            with open(fpath, 'w', encoding='utf-8') as f:
+                json.dump(check_json_format(self.__dict__), f, ensure_ascii=False, indent=2)
+            config_file = os.getenv('SWIFT_CONFIG_FILE')
+            if config_file:
+                shutil.copy(config_file, output_dir)
+
+    def _init_device(self):
+        if is_dist():
+            set_device()
+
+    def get_template(self, processor: Optional[Processor] = None, **kwargs) -> Template:
+        if processor is None:
+            processor = self.get_model_processor(load_model=False)[1]
+        template_kwargs = self.get_template_kwargs()
+        if 'template_type' in kwargs:
+            template_type = kwargs.get('template_type')
+        else:
+            template_type = self.template
+        template_kwargs['template_type'] = template_type
+        template = get_template(processor, **template_kwargs)
+        return template
+
+    def get_model_processor(self,
+                            *,
+                            model=None,
+                            model_type=None,
+                            revision=None,
+                            task_type=None,
+                            num_labels=None,
+                            **kwargs):
+        if self.tuner_backend == 'unsloth':
+            return load_by_unsloth(self)
+        res = self.get_model_kwargs()
+        res.update(kwargs)
+        # compat rlhf
+        res['model_id_or_path'] = model or self.model
+        res['model_type'] = model_type or self.model_type
+        res['revision'] = revision or self.model_revision
+        res['task_type'] = task_type or self.task_type
+        res['num_labels'] = num_labels or self.num_labels
+
+        return get_model_processor(**res)
+
+    def load_dataset(self):
+        dataset_kwargs = self.get_dataset_kwargs()
+        train_dataset, val_dataset = None, None
+        if self.dataset:
+            train_dataset, val_dataset = load_dataset(
+                self.dataset,
+                split_dataset_ratio=self.split_dataset_ratio,
+                shuffle=self.dataset_shuffle,
+                **dataset_kwargs)
+        if len(self.val_dataset) > 0:
+            # Loading val dataset
+            dataset_kwargs.pop('interleave_prob', None)
+            _, val_dataset = load_dataset(
+                self.val_dataset, split_dataset_ratio=1.0, shuffle=self.val_dataset_shuffle, **dataset_kwargs)
+            assert self.split_dataset_ratio == 0.
+        return train_dataset, val_dataset

swift/arguments/base_args/data_args.py

@@ -0,0 +1,145 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from dataclasses import dataclass, field
+from typing import List, Literal, Optional, Union
+
+from swift.dataset import register_dataset_info
+from swift.utils import get_logger, json_parse_to_dict
+
+logger = get_logger()
+
+
+@dataclass
+class DataArguments:
+    """Holds arguments related to dataset handling and processing.
+
+    Args:
+        dataset (List[str]): A list of dataset IDs or paths. Defaults to [].
+            Format for each dataset: 'dataset_id_or_path:subset#count'. Both subset and count are optional.
+            - Subsets: Only effective for dataset IDs or folders. Use '/' to select multiple subsets (e.g.,
+            'dataset_id:subset1/subset2') or 'all' to select all registered subsets. If only one subset is
+            registered, it will be used by default; otherwise, 'default' is the default.
+            - Sampling Count: By default, the full dataset is used. Use '#count' to sample. If count <
+            total samples, it performs random sampling without replacement. If count > total, it repeats
+            the full dataset `count // total` times and then randomly samples an additional `count % total`
+            samples. Note: Streaming datasets or setting `--dataset_shuffle false` will result in sequential
+            sampling.
+            - Local datasets: Supports formats like jsonl, csv, json, and folders.
+        val_dataset (List[str]): A list of validation dataset IDs or paths. Defaults to [].
+        cached_dataset (List[str]): Use cached datasets to avoid GPU time being occupied by tokenization during
+            training/inference on large datasets. This parameter is used to set the folder path(s) of
+            cached training datasets, and defaults to `[]`.
+            This is generated by the `swift export --to_cached_dataset true ...` command.
+            ms-swift only stores an extra 'length' field and filters out erroneous samples
+            to reduce storage. Actual preprocessing happens concurrently with training.
+        cached_val_dataset (List[str]): Folder path(s) for cached validation datasets, default is [].
+        split_dataset_ratio (float): The ratio to split from the training set for validation if `val_dataset` is not
+            provided. Defaults to 0.0. Note: The default was 0.01 in `ms-swift<3.6`.
+        data_seed (int): The random seed for dataset shuffling. Defaults to 42.
+        dataset_num_proc (int): The number of processes to use for dataset preprocessing. Defaults to 1.
+        load_from_cache_file (bool): Whether to load the dataset from cache files. Recommended to set to `True` during
+            actual runs and `False` during debugging. Defaults to False.
+            Note: The default was `True` in `ms-swift<3.9`.
+        dataset_shuffle (bool): Whether to shuffle the training dataset. Defaults to True.
+            Note: For CPT/SFT, shuffling occurs at both the dataset level (controlled by this flag) and the dataloader
+            level.
+        val_dataset_shuffle (bool): Whether to shuffle the validation dataset. Defaults to False.
+        streaming (bool): Enables streaming to read and process the dataset on-the-fly. `--max_steps` must be set as the
+            dataset length is unknown. This allows preprocessing to overlap with training but can become a bottleneck
+            with a large `world_size` as preprocessing only runs on rank 0. Defaults to False.
+        interleave_prob (Optional[List[float]]): If set, combines datasets using `interleave_datasets` with the
+            provided probabilities instead of `concatenate_datasets`. Typically used for streaming. Defaults to None.
+        stopping_strategy (str): The stopping strategy for `interleave_datasets`. Can be "first_exhausted" or
+            "all_exhausted". Defaults to "first_exhausted".
+        shuffle_buffer_size (int): The buffer size for shuffling in streaming mode. Only effective if `dataset_shuffle`
+            is `True`. Defaults to 1000.
+        download_mode (str): The dataset download mode. Options are 'reuse_dataset_if_exists' and 'force_redownload'.
+            Defaults to 'reuse_dataset_if_exists'.
+        columns (Optional[str]): A JSON string for column mapping to fit the format required by `AutoPreprocessor`.
+            Example: '{"text1": "query", "text2": "response"}'. Defaults to None.
+        strict (bool): If `True`, raises an error on any problematic data row. If `False`, discards the problematic
+            sample and continues. Typically used for debugging. Defaults to False.
+        remove_unused_columns (bool): Whether to remove columns not used by the model. If `False`, extra columns are
+            passed to the trainer's `compute_loss` function, which is useful for custom loss calculations.
+            Defaults to True. Note: The default is `False` for GPRO.
+        disable_auto_column_mapping (bool): By default, column names in the dataset are automatically mapped.
+            This parameter disables that behavior (the `columns` parameter remains effective), defaulting to `False`.
+        model_name (Optional[List[str]]): For self-cognition tasks, replaces the `{{NAME}}` placeholder in the
+            `swift/self-cognition` dataset. Pass Chinese and English names.
+            Example: `--model_name 小黄 'Xiao Huang'`. Defaults to None.
+        model_author (Optional[List[str]]): For self-cognition tasks, replaces the `{{AUTHOR}}` placeholder in the
+            `swift/self-cognition` dataset. Pass author's Chinese and English names.
+            Example: `--model_author '魔搭' 'ModelScope'`. Defaults to None.
+        custom_dataset_info (List[str]): Path to a custom dataset registration JSON file. Defaults to [].
+    """
+    # dataset_id or dataset_dir or dataset_path
+    dataset: List[str] = field(default_factory=list)
+    val_dataset: List[str] = field(default_factory=list)
+    cached_dataset: List[str] = field(default_factory=list)
+    cached_val_dataset: List[str] = field(default_factory=list)
+    split_dataset_ratio: float = 0.
+
+    data_seed: int = 42
+    dataset_num_proc: int = 1
+    load_from_cache_file: bool = False
+    dataset_shuffle: bool = True
+    val_dataset_shuffle: bool = False
+    streaming: bool = False
+    interleave_prob: Optional[List[float]] = None
+    stopping_strategy: Literal['first_exhausted', 'all_exhausted'] = 'first_exhausted'
+    shuffle_buffer_size: int = 1000
+
+    download_mode: Literal['force_redownload', 'reuse_dataset_if_exists'] = 'reuse_dataset_if_exists'
+    columns: Optional[Union[dict, str]] = None
+    strict: bool = False
+    remove_unused_columns: bool = True
+    disable_auto_column_mapping: bool = False
+    # Chinese name and English name
+    model_name: Optional[List[str]] = field(default=None, metadata={'help': "e.g. ['小黄', 'Xiao Huang']"})
+    model_author: Optional[List[str]] = field(default=None, metadata={'help': "e.g. ['魔搭', 'ModelScope']"})
+
+    custom_dataset_info: List[str] = field(default_factory=list)  # .json
+
+    def _init_custom_dataset_info(self):
+        """register custom dataset_info.json to datasets"""
+        if isinstance(self.custom_dataset_info, str):
+            self.custom_dataset_info = [self.custom_dataset_info]
+        for path in self.custom_dataset_info:
+            register_dataset_info(path)
+
+    def __post_init__(self):
+        self.columns = json_parse_to_dict(self.columns)
+        if len(self.val_dataset) > 0 or self.streaming and self.split_dataset_ratio > 0:
+            self.split_dataset_ratio = 0.
+            if len(self.val_dataset) > 0:
+                msg = 'len(args.val_dataset) > 0'
+            else:
+                msg = 'args.streaming is True'
+            logger.info(f'Because {msg}, setting split_dataset_ratio: {self.split_dataset_ratio}')
+        self._init_custom_dataset_info()
+        if isinstance(self.cached_dataset, str):
+            self.cached_dataset = [self.cached_dataset]
+        self._init_val_dataset_exists()
+
+    def _init_val_dataset_exists(self):
+        self._val_dataset_exists = bool(self.dataset and self.split_dataset_ratio > 0 or self.val_dataset
+                                        or self.cached_val_dataset)
+
+    def get_dataset_kwargs(self):
+        return {
+            'seed': self.data_seed,
+            'num_proc': self.dataset_num_proc,
+            'load_from_cache_file': self.load_from_cache_file,
+            'streaming': self.streaming,
+            'interleave_prob': self.interleave_prob,
+            'stopping_strategy': self.stopping_strategy,
+            'shuffle_buffer_size': self.shuffle_buffer_size,
+            'use_hf': self.use_hf,
+            'hub_token': self.hub_token,
+            'download_mode': self.download_mode,
+            'columns': self.columns,
+            'strict': self.strict,
+            'model_name': self.model_name,
+            'model_author': self.model_author,
+            'remove_unused_columns': self.remove_unused_columns,
+            'disable_auto_column_mapping': self.disable_auto_column_mapping,
+        }

swift/arguments/base_args/generation_args.py

@@ -0,0 +1,76 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+from swift.infer_engine import RequestConfig
+from swift.utils import get_logger
+
+logger = get_logger()
+
+
+@dataclass
+class GenerationArguments:
+    """A dataclass that holds arguments for text generation.
+
+    Args:
+        max_new_tokens (Optional[int]): The maximum number of new tokens to generate. Defaults to None (unlimited).
+        temperature (Optional[float]): The sampling temperature. A higher temperature makes the output more random. To
+            disable randomness, you can set this to 0 or `top_k` to 1. Defaults to None, which means loading from
+            'generation_config.json'.
+        top_k (Optional[int]): The number of highest probability vocabulary tokens to keep for top-k-filtering.
+            Defaults to None (reads from 'generation_config.json').
+        top_p (Optional[float]): The cumulative probability for nucleus sampling. Filters the vocabulary to the
+            smallest set of tokens whose cumulative probability exceeds `top_p`. Defaults to None (reads from
+            'generation_config.json').
+        repetition_penalty (Optional[float]): The penalty applied to repeated tokens. A value of 1.0 means no penalty.
+            Defaults to None (reads from 'generation_config.json').
+        num_beams (Optional[int]): The number of beams to use for beam search. Defaults to 1.
+        stream (bool): Whether to enable streaming output. Defaults to None, which is `True` for interactive mode and
+            `False` for batch inference. Note: For ms-swift < 3.6, the default is `False`.
+        stop_words (List[str]): A list of extra stop words, in addition to the end-of-sequence token. Note: The
+            `eos_token` is removed from the output, while these stop words are preserved. Defaults to an empty list.
+        logprobs (bool): Whether to output log probabilities of the generated tokens. Defaults to False.
+        top_logprobs (Optional[int]): The number of top log probabilities to return for each token position. Requires
+            `logprobs` to be True. Defaults to None.
+        structured_outputs_regex (Optional[str]): A regular expression pattern for structured outputs (guided decoding).
+            When set, the model's generation is constrained to match the specified regex pattern. This is useful for
+            tasks requiring structured outputs like reasoning chains. Only effective when `infer_backend` is 'vllm'.
+            Defaults to None.
+    """
+
+    # generation config
+    max_new_tokens: Optional[int] = None  # Unlimited, constrained by max_model_len.
+    # If it is None, use the parameters from generation_config.
+    temperature: Optional[float] = None  # Set to 0, which means do_sample is False.
+    top_k: Optional[int] = None
+    top_p: Optional[float] = None
+    repetition_penalty: Optional[float] = None
+    num_beams: int = 1
+
+    stream: Optional[bool] = None
+    stop_words: List[str] = field(default_factory=list)
+    logprobs: bool = False
+    top_logprobs: Optional[int] = None
+    # structured outputs (guided decoding), only effective for vllm backend
+    structured_outputs_regex: Optional[str] = None
+
+    def _init_stream(self):
+        if self.stream is None:
+            self.stream = False
+
+    def get_request_config(self):
+        if getattr(self, 'task_type') != 'causal_lm':
+            return
+
+        return RequestConfig(
+            max_tokens=self.max_new_tokens,
+            temperature=self.temperature,
+            top_p=self.top_p,
+            top_k=self.top_k,
+            num_beams=self.num_beams,
+            stop=self.stop_words,
+            stream=self.stream,
+            repetition_penalty=self.repetition_penalty,
+            logprobs=self.logprobs,
+            top_logprobs=self.top_logprobs,
+            structured_outputs_regex=self.structured_outputs_regex)

swift/arguments/base_args/model_args.py

@@ -0,0 +1,249 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import ast
+import math
+import os
+import torch
+from dataclasses import dataclass, field
+from transformers.utils import is_torch_mps_available
+from typing import Any, Dict, List, Literal, Optional, Union
+
+from swift.model import MODEL_MAPPING, get_model_info_meta, get_model_name
+from swift.utils import HfConfigFactory, get_dist_setting, get_logger, json_parse_to_dict
+
+logger = get_logger()
+
+
+@dataclass
+class ModelArguments:
+    """A dataclass that holds various arguments related to model configuration and usage.
+
+    Args:
+        model (Optional[str]): The model ID from the Hub or a local path to the model. Defaults to None.
+        model_type (Optional[str]): The model type. In ms-swift, a 'model_type' groups models with the same
+            architecture, loading process, and template. Defaults to None, which enables auto-selection based on
+            the suffix of `--model` and the 'architectures' attribute in `config.json`. The `model_type` for a
+            corresponding model can be found in the list of supported models. Note: The concept of `model_type`
+            in ms-swift differs from the `model_type` in `config.json`. Custom models usually require registering
+            their own `model_type` and `template`.
+        model_revision (Optional[str]): The revision of the model. Defaults to None.
+        task_type (str): The task type. Can be 'causal_lm', 'seq_cls', 'embedding', 'reranker', or
+            'generative_reranker'. If set to 'seq_cls', you usually need to specify `--num_labels` and
+            `--problem_type`. Defaults to 'causal_lm'.
+        torch_dtype (Optional[str]): The data type of the model weights. Supports 'float16', 'bfloat16', 'float32'.
+            Defaults to None, in which case it's read from the 'config.json' file.
+        attn_impl (Optional[str]): The attention implementation to use. Options include 'sdpa', 'eager', 'flash_attn',
+            'flash_attention_2', 'flash_attention_3', 'flash_attention_4', etc.
+            Defaults to None, which means it will be read from 'config.json'.
+            Note: Support for these implementations depends on the model's transformers implementation.
+            If set to 'flash_attn' (for backward compatibility), 'flash_attention_2' will be used.
+        experts_impl (Optional[str]): Expert implementation type, options are 'grouped_mm', 'batched_mm', 'eager'.
+            Defaults to None. This feature requires "transformers>=5.0.0".
+        new_special_tokens (List[str]): Additional special tokens to be added to the tokenizer. Can also be a path to
+            a `.txt` file, where each line is a special token. Defaults to an empty list `[]`.
+        num_labels (Optional[int]): The number of labels for classification tasks (when `--task_type` is 'seq_cls').
+            Required for such tasks. Defaults to None.
+        problem_type (Optional[str]): The problem type for classification tasks (`--task_type` 'seq_cls'). Options are
+            'regression', 'single_label_classification', 'multi_label_classification'. Defaults to None, but is
+            automatically set to 'regression' if the model is a reward_model or `num_labels` is 1, and
+            'single_label_classification' otherwise.
+        rope_scaling (Optional[str]): The RoPE scaling type. You can pass a string like 'linear', 'dynamic', or
+            'yarn', and ms-swift will automatically set the corresponding `rope_scaling` and override the
+            'config.json' value. Alternatively, you can pass a JSON string (e.g., '{"factor":2.0, "type":"yarn"}'),
+            which will directly override the `rope_scaling` in 'config.json'. Defaults to None.
+        device_map (Optional[str]): The device map configuration for the model, e.g., 'auto', 'cpu', a JSON string,
+            or a path to a JSON file. This argument is passed directly to the `from_pretrained` method of transformers.
+            Defaults to None, and will be set automatically based on the device and distributed training settings.
+        max_memory (Optional[str]): The maximum memory allocation for each device when `device_map` is 'auto' or
+            'sequential'. Example: '{0: "20GB", 1: "20GB"}'. This argument is passed directly to the `from_pretrained`
+            method of transformers. Defaults to None.
+        max_model_len (Optional[int]): The maximum model length. This is used to calculate the RoPE scaling factor
+            when `rope_scaling` is specified as a string. If not None, it overrides the `max_position_embeddings`
+            value in 'config.json'. Defaults to None.
+        local_repo_path (Optional[str]): Path to a local repository for models that require a GitHub repo during
+            loading (e.g., deepseek-vl2). This avoids network issues during `git clone`. Defaults to None.
+        init_strategy (Optional[str]): The strategy to initialize all uninitialized parameters when loading a model
+            (especially for custom architectures). Options include 'zero', 'uniform', 'normal', 'xavier_uniform',
+            'xavier_normal', 'kaiming_uniform', 'kaiming_normal', 'orthogonal'. Defaults to None.
+    """
+    model: Optional[str] = None  # model id or model path
+    model_type: Optional[str] = field(
+        default=None, metadata={'help': f'model_type choices: {list(MODEL_MAPPING.keys())}'})
+    model_revision: Optional[str] = None
+    task_type: Literal['causal_lm', 'seq_cls', 'embedding', 'reranker', 'generative_reranker'] = None
+
+    torch_dtype: Literal['bfloat16', 'float16', 'float32', None] = None
+    # flash_attn: It will automatically convert names based on the model.
+    # None: It will be automatically selected between sdpa and eager.
+    # 'flash_attn', 'sdpa', 'eager', 'flex_attention',
+    # 'flash_attention_2', 'flash_attention_3', 'flash_attention_4'
+    attn_impl: Optional[str] = None
+    experts_impl: Optional[str] = None
+    new_special_tokens: List[str] = field(default_factory=list)
+
+    num_labels: Optional[int] = None
+    problem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] = None
+    rope_scaling: Optional[str] = None
+    device_map: Optional[Union[dict, str]] = None
+    max_memory: Optional[Union[dict, str]] = None
+    max_model_len: Optional[int] = None
+    # When some model code needs to be downloaded from GitHub,
+    # this parameter specifies the path to the locally downloaded repository.
+    local_repo_path: Optional[str] = None
+    init_strategy: Literal['zero', 'uniform', 'normal', 'xavier_uniform', 'xavier_normal', 'kaiming_uniform',
+                           'kaiming_normal', 'orthogonal'] = None
+
+    def _init_device_map(self):
+        """Prepare device map args"""
+        if self.device_map:
+            self.device_map: Union[str, Dict[str, Any], None] = json_parse_to_dict(self.device_map, strict=False)
+        # compat mp&ddp
+        _, local_rank, _, local_world_size = get_dist_setting()
+        if local_world_size > 1 and isinstance(self.device_map, dict) and local_rank > 0:
+            for k, v in self.device_map.items():
+                if isinstance(v, int):
+                    self.device_map[k] += local_rank
+
+    def _init_max_memory(self):
+        if isinstance(self.max_memory, str):
+            try:
+                self.max_memory = ast.literal_eval(self.max_memory)
+            except Exception:
+                pass
+        self.max_memory = json_parse_to_dict(self.max_memory)
+        # compat mp&ddp
+        _, local_rank, _, local_world_size = get_dist_setting()
+        if local_world_size > 1 and isinstance(self.max_memory, dict) and local_rank > 0:
+            for k in list(self.max_memory.keys()):
+                if isinstance(k, int):
+                    self.max_memory[k + local_rank] = self.max_memory.pop(k)
+
+    def _init_torch_dtype(self) -> None:
+        """"If torch_dtype is None, find a proper dtype by the config.json/GPU"""
+        from ..sft_args import SftArguments
+
+        self.torch_dtype: Optional[torch.dtype] = HfConfigFactory.to_torch_dtype(self.torch_dtype)
+        self.torch_dtype: torch.dtype = self._init_model_info()
+        # Mixed Precision Training
+        if isinstance(self, SftArguments):
+            self._init_mixed_precision()
+
+    def _init_mixed_precision(self):
+        if is_torch_mps_available():
+            fp16, bf16 = False, False
+        elif self.torch_dtype in {torch.float16, torch.float32}:
+            fp16, bf16 = True, False
+        elif self.torch_dtype == torch.bfloat16:
+            fp16, bf16 = False, True
+        else:
+            raise ValueError(f'args.torch_dtype: {self.torch_dtype}')
+        if self.fp16 is None:
+            self.fp16 = fp16
+        if self.bf16 is None:
+            self.bf16 = bf16
+
+    def _init_rope_scaling(self):
+        if self.rope_scaling:
+            rope_scaling: dict = json_parse_to_dict(self.rope_scaling, strict=False)
+            if isinstance(rope_scaling, str):
+                assert rope_scaling in ['linear', 'dynamic', 'yarn']
+                rope_scaling = {'type': rope_scaling}
+        else:
+            rope_scaling = self.model_info.rope_scaling
+            # reset the factor
+            rope_scaling.pop('factor', None)
+
+        rope_type = rope_scaling.get('rope_type', rope_scaling.get('type', 'default'))
+        if 'factor' not in rope_scaling and self.max_model_len is None and rope_type == 'default':
+            # fix megatron qwen2_5_vl
+            self.rope_scaling = rope_scaling
+            logger.info(f'Setting args.rope_scaling: {rope_scaling}')
+            return
+
+        # get origin_max_model_len
+        origin_max_model_len = None
+        if rope_scaling and rope_scaling.get('original_max_position_embeddings') is not None:
+            origin_max_model_len = rope_scaling['original_max_position_embeddings']
+        elif self.model_info.rope_scaling:
+            if self.model_info.rope_scaling.get('original_max_position_embeddings') is not None:
+                origin_max_model_len = self.model_info.rope_scaling['original_max_position_embeddings']
+            elif self.model_info.rope_scaling.get('factor') is not None:
+                origin_max_model_len = self.model_info.max_model_len // self.model_info.rope_scaling['factor']
+        if origin_max_model_len is None:
+            origin_max_model_len = self.model_info.max_model_len
+        assert origin_max_model_len is not None, '`origin_max_model_len` from model config is not set'
+        rope_scaling['original_max_position_embeddings'] = origin_max_model_len
+
+        if 'factor' not in rope_scaling:
+            assert self.max_model_len is not None, (
+                'max_model_len must be set if rope_scaling does not contain a "factor"')
+            rope_scaling['factor'] = max(float(math.ceil(self.max_model_len / origin_max_model_len)), 1.0)
+        rope_model_len = int(origin_max_model_len * rope_scaling['factor'])
+        if self.max_model_len is None:
+            self.max_model_len = rope_model_len
+        elif self.max_model_len > rope_model_len:
+            logger.warning(f'rope config ({rope_model_len} = {rope_scaling["factor"]} * '
+                           f'{origin_max_model_len}) should be bigger than max_model_len '
+                           f'from command line ({self.max_model_len})')
+        self.rope_scaling = rope_scaling
+        logger.info(f'Setting args.rope_scaling: {rope_scaling}')
+        logger.info(f'Setting args.max_model_len: {self.max_model_len}')
+
+    def _init_model_info(self) -> torch.dtype:
+        model_kwargs = self.get_model_kwargs()
+        if self.tuner_backend == 'unsloth':
+            model_kwargs['download_model'] = True
+        self.model_info, self.model_meta = get_model_info_meta(**model_kwargs)
+        self.task_type = self.model_info.task_type
+        self.num_labels = self.model_info.num_labels
+
+        self.model_dir = self.model_info.model_dir
+        self.model_type = self.model_info.model_type
+        if self.rope_scaling or self.model_info.rope_scaling and self.max_model_len is not None:
+            self._init_rope_scaling()
+        return self.model_info.torch_dtype
+
+    def _init_new_special_tokens(self):
+        if isinstance(self.new_special_tokens, str):
+            self.new_special_tokens = [self.new_special_tokens]
+        new_special_tokens = []
+        for token in self.new_special_tokens:
+            if token.endswith('.txt'):
+                assert os.path.isfile(token), f'special_tokens_path: {token}'
+                with open(token, 'r', encoding='utf-8') as f:
+                    text = f.read()
+                new_special_tokens += text.split()
+            else:
+                new_special_tokens.append(token)
+        self.new_special_tokens = new_special_tokens
+
+    def __post_init__(self):
+        if self.model is None:
+            raise ValueError(f'Please set --model <model_id_or_path>`, model: {self.model}')
+        self._init_new_special_tokens()
+        self.model_suffix = get_model_name(self.model)
+        self._init_device_map()
+        self._init_max_memory()
+        self._init_torch_dtype()
+
+    def get_model_kwargs(self):
+        return {
+            'model_id_or_path': self.model,
+            'torch_dtype': self.torch_dtype,
+            'model_type': self.model_type,
+            'revision': self.model_revision,
+            'use_hf': self.use_hf,
+            'hub_token': self.hub_token,
+            'local_repo_path': self.local_repo_path,
+            'device_map': self.device_map,
+            'max_memory': self.max_memory,
+            'quantization_config': self.get_quantization_config(),
+            'attn_impl': self.attn_impl,
+            'experts_impl': self.experts_impl,
+            'new_special_tokens': self.new_special_tokens,
+            'rope_scaling': self.rope_scaling,
+            'max_model_len': self.max_model_len,
+            'task_type': self.task_type,
+            'num_labels': self.num_labels,
+            'problem_type': self.problem_type,
+            'init_strategy': self.init_strategy,
+        }

swift/arguments/base_args/quant_args.py

@@ -0,0 +1,122 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import torch
+from dataclasses import dataclass
+from typing import Literal, Optional
+
+from swift.model import get_model_processor
+from swift.utils import HfConfigFactory, get_modules_to_not_convert
+
+
+@dataclass
+class QuantizeArguments:
+    """A dataclass that holds the configuration for model quantization.
+
+    Args:
+        quant_method (Optional[str]): The quantization method to use when loading the model. Can be one of {'bnb',
+            'hqq', 'eetq', 'quanto', 'fp8'}. Note: This is not required for QLoRA training on pre-quantized AWQ/GPTQ
+            models. Defaults to None.
+        quant_bits (Optional[Union[int, str]]): The number of bits for quantization, e.g., {1, 2, 3, 4, 8, 'float8'}.
+            Defaults to None.
+        hqq_axis (Optional[int]): The quantization axis for HQQ quantization. Defaults to None.
+        bnb_4bit_compute_dtype (Optional[str]): The compute data type for 4-bit BNB quantization. Can be one of {
+            'float16', 'bfloat16', 'float32'}. Defaults to None, which will use the model's `torch_dtype`.
+        bnb_4bit_quant_type (str): The quantization type for 4-bit BNB quantization. Can be one of {'fp4', 'nf4'}.
+            Defaults to 'nf4'.
+        bnb_4bit_use_double_quant (bool): Whether to use double quantization for 4-bit BNB quantization.
+            Defaults to True.
+        bnb_4bit_quant_storage (Optional[str]): The storage type for packing quantized 4-bit parameters in BNB.
+            Defaults to None.
+    """
+    # awq, gptq, and aqlm need to be pre-quantized models.
+    #   It can be detected automatically, without the need to pass in.
+    # while bnb, hqq, and eetq can be quantized during SFT using the original models.
+    quant_method: Literal['bnb', 'hqq', 'eetq', 'quanto', 'fp8'] = None
+    # bnb: 4,8; hqq: 1,2,3,4,8'; eetq: 8
+    # awq: 4; gptq: 2,3,4,8
+    quant_bits: Literal[1, 2, 3, 4, 8, 'float8'] = None
+    # hqq
+    hqq_axis: Optional[int] = None
+    # bnb
+    bnb_4bit_compute_dtype: Literal['float16', 'bfloat16', 'float32', None] = None
+    bnb_4bit_quant_type: Literal['fp4', 'nf4'] = 'nf4'
+    bnb_4bit_use_double_quant: bool = True
+    bnb_4bit_quant_storage: Optional[str] = None
+
+    def get_quantization_config(self):
+        if self.quant_method is None or self.quant_method in {'awq', 'gptq', 'gptq_v2'}:
+            return None
+        assert self.quant_method in {'bnb', 'hqq', 'eetq', 'quanto', 'fp8'}
+        if self.quant_method != 'fp8' and self.quant_bits is None:
+            raise ValueError(f'Please set the quant_bits. args.quant_bits: {self.quant_bits}')
+        if self.quant_method == 'bnb':
+            if self.quant_bits == 4:
+                load_in_4bit, load_in_8bit = True, False
+            elif self.quant_bits == 8:
+                load_in_4bit, load_in_8bit = False, True
+            else:
+                raise ValueError(f'bnb not support quant_bits: {self.quant_bits}')
+
+            from transformers import BitsAndBytesConfig
+            llm_int8_skip_modules = self.get_modules_to_not_convert()
+            quantization_config = BitsAndBytesConfig(
+                load_in_4bit=load_in_4bit,
+                load_in_8bit=load_in_8bit,
+                bnb_4bit_compute_dtype=self.bnb_4bit_compute_dtype,
+                bnb_4bit_quant_type=self.bnb_4bit_quant_type,
+                bnb_4bit_use_double_quant=self.bnb_4bit_use_double_quant,
+                bnb_4bit_quant_storage=self.bnb_4bit_quant_storage,
+                llm_int8_skip_modules=llm_int8_skip_modules)
+        elif self.quant_method == 'fp8':
+            if not hasattr(self, 'model_info'):
+                return
+            from transformers import FineGrainedFP8Config
+            with torch.device('meta'):
+                hf_model, _ = get_model_processor(self.model_dir, model_type=self.model_type, return_dummy_model=True)
+            modules_to_not_convert = get_modules_to_not_convert(hf_model)
+            quantization_config = FineGrainedFP8Config(modules_to_not_convert=modules_to_not_convert)
+        elif self.quant_method == 'hqq':
+            from transformers import HqqConfig
+            quantization_config = HqqConfig(nbits=self.quant_bits, axis=self.hqq_axis)
+        elif self.quant_method == 'quanto':
+            from transformers import QuantoConfig
+            if self.quant_bits == 8:
+                weights = 'int8'
+            elif self.quant_bits == 'float8':
+                weights = 'float8'
+            elif self.quant_bits == 4:
+                weights = 'int4'
+            elif self.quant_bits == 2:
+                weights = 'int2'
+            else:
+                raise ValueError('quanto quantization only support quant bits 2/4/8/float8')
+            quantization_config = QuantoConfig(weights=weights)
+        else:  # 'eetq'
+            from transformers import EetqConfig
+            quantization_config = EetqConfig(f'int{self.quant_bits}')
+
+        return quantization_config
+
+    def get_modules_to_not_convert(self):
+        if not hasattr(self, 'model_meta') or not hasattr(self, 'model_info'):
+            return None
+        model_arch = self.model_meta.model_arch
+        res = []
+        if self.model_info.is_moe_model:
+            res += ['mlp.gate', 'mlp.shared_expert_gate']
+        if model_arch is not None:
+            for key in ['vision_tower', 'aligner']:
+                value = getattr(model_arch, key, None)
+                if value:
+                    res += value
+        if not res:
+            return None
+        res.append('lm_head')
+        return res
+
+    def __post_init__(self):
+        if self.bnb_4bit_compute_dtype is None:
+            if self.torch_dtype in {torch.float16, torch.float32}:
+                self.bnb_4bit_compute_dtype = torch.float32
+            elif self.torch_dtype == torch.bfloat16:
+                self.bnb_4bit_compute_dtype = torch.bfloat16
+        self.bnb_4bit_compute_dtype: torch.dtype = HfConfigFactory.to_torch_dtype(self.bnb_4bit_compute_dtype)

swift/arguments/base_args/template_args.py

@@ -0,0 +1,204 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import os
+from dataclasses import dataclass, field
+from typing import Literal, Optional
+
+from swift.template import TEMPLATE_MAPPING, get_template_meta
+from swift.utils import get_logger
+
+logger = get_logger()
+
+
+@dataclass
+class TemplateArguments:
+    """TemplateArguments class holds various arguments for template configuration.
+
+    This dataclass manages settings related to how data is formatted and processed using templates, including
+    tokenization, truncation, loss calculation, and special handling for multimodal and agent-based models.
+
+    Args:
+        template (Optional[str]): The dialogue template type. Defaults to None, which automatically selects the
+            template corresponding to the model type. Refer to the list of supported models for mappings.
+        system (Optional[str]): Custom system prompt. Can be a string or a path to a .txt file. Defaults to None,
+            which uses the default system from the registered template.
+            Note: The priority for the system prompt is as follows:
+            1. System prompt from the dataset.
+            2. The `--system` command-line argument.
+            3. The `default_system` set when the template was registered.
+        max_length (Optional[int]): The maximum number of tokens for a single sample after tokenization. Samples
+            exceeding this length are handled according to `truncation_strategy` to prevent OOM errors. Defaults to
+            None, which uses the model's maximum supported length (`max_model_len`). In PPO, GRPO, and inference
+            scenarios, this argument specifies the `max_prompt_length`.
+        truncation_strategy (Literal['delete', 'left', 'right', 'split']): Strategy for handling samples exceeding
+            `max_length`. Options are 'delete', 'left' (truncate from the left), 'right' (truncate from the right),
+            and 'split' (split into multiple samples). Defaults to 'delete'.
+            Note: The 'split' strategy is only supported during pre-training (e.g., `swift/megatron pt`),
+            and is incompatible with `cached_dataset`. It splits long samples to avoid wasting tokens.
+            Note: For multimodal models, setting this to 'left' or 'right' preserves all image tokens, which may lead
+            to OOM errors.
+        max_pixels (Optional[int]): The maximum number of pixels (H*W) for an input image in a multimodal model.
+            Images exceeding this limit will be scaled down to prevent OOM errors. Defaults to None, meaning no limit.
+            Note: This parameter applies to all multimodal models. The model-specific `MAX_PIXELS` parameter for
+            Qwen2.5-VL is separate and only applies to that model.
+        agent_template (Optional[str]): The Agent template to use. This determines how the 'tools' list is converted
+            into a 'system' prompt, how tool calls are extracted from the model's response during inference, and the
+            format for tool call messages. Options include "react_en", "hermes", "glm4", "qwen_en", "toolbench", etc.
+            Defaults to None, which auto-selects based on the model type. Refer to the Agent documentation for more
+            details.
+        norm_bbox (Optional[Literal['norm1000', 'none']]): Controls how bounding box coordinates (from the "bbox"
+            field in the dataset) are scaled. 'norm1000' scales coordinates to a 1000x1000 grid, while 'none' performs
+            no scaling. Defaults to None, which auto-selects based on the model. This handles cases where images are
+            resized during training (e.g., due to `max_pixels`).
+        use_chat_template (bool): Whether to use the chat template or the generation template. The generation template
+            is typically used for pre-training. Defaults to True.
+            Note: Defaults to False for `swift pt`, which uses the generation template. This parameter is compatible
+            with multimodal models.
+        padding_side (Literal['left', 'right']): The side to pad on when `batch_size >= 2` during training.
+            Options are 'left' or 'right'. Defaults to 'right'. For inference with `batch_size >= 2`, padding is always
+            on the left.
+            Note: Defaults to 'left' for PPO and GKD.
+        padding_free (bool): If True, flattens the data within a batch to avoid padding, reducing memory usage and
+            speeding up training. Sequences within the batch remain causally isolated. Defaults to False. Supported for
+            CPT/SFT/DPO/GRPO/KTO/GKD.
+            Note: This requires `--attn_impl flash_attn` and `transformers>=4.44`. Compared to packing, padding_free
+            has no preprocessing overhead, but packing offers faster training speeds and more stable memory usage.
+        loss_scale (str): Loss weight configuration for training tokens. Default is `'default'`.
+            loss_scale includes 3 basic strategies: 'default', 'last_round', 'all', and other strategies:
+            'ignore_empty_think' and agent-specific ones: 'react', 'hermes', 'qwen', 'agentflan', 'alpha_umi', etc.
+            For available options, refer to
+            [loss_scale module](https://github.com/modelscope/ms-swift/blob/main/swift/loss_scale/mapping.py).
+            ms-swift supports mixing basic strategies with other strategies,
+            for example: `'default+ignore_empty_think'`, `'last_round+ignore_empty_think'`.
+            If no basic strategy is specified, it defaults to 'default',
+            for example: 'hermes' is equivalent to 'default+hermes'.
+            Multiple non-base strategies can be chained together
+            (each strategy processes the output segments of the previous one, with weights
+            multiplied accordingly). For example: `'last_round+hermes+ignore_empty_think'`, where
+            `'last_round'` is the base strategy, and `'hermes+ignore_empty_think'` represents a
+            chain of multiple non-base strategies that share the same base strategy.
+            - 'default': All responses (including history) are calculated with weight 1 for cross-entropy loss
+            (**system/user/multimodal tokens in messages and `tool_response` parts in Agent training are
+            not included in loss calculation**). (**Default value for SFT**)
+            - 'last_round': Only calculate loss for the last round response. The last round
+            means all content after the last "user". (**Default value for RLHF**)
+            - 'all': Calculate loss for all tokens. (**Default value for `swift pt`**)
+            - 'ignore_empty_think': Ignore loss computation for empty `'<think>\n\n</think>\n\n'`
+            (as long as it matches the regex `'<think>\\s*</think>\\s*'`).
+            - 'react', 'hermes', 'qwen': Adjust the loss weight of the `tool_call` part to 2.
+        sequence_parallel_size (int): The size of sequence parallelism. Defaults to 1. Currently supported for CPT,
+            SFT, DPO, and GRPO.
+        template_backend (Literal['swift', 'jinja']): The backend to use for templating. Options are 'swift' or
+            'jinja'. Defaults to 'swift'. If 'jinja' is used, it will leverage `transformers.apply_chat_template`.
+            Note: The 'jinja' backend is only supported for inference, not for training, as it cannot determine the
+            token range for loss calculation.
+        response_prefix (Optional[str]): A prefix string for the response, e.g., '<think>\\n' for Qwen-32B. This
+            parameter only affects inference. Defaults to None, which is auto-set based on the model.
+        enable_thinking (Optional[bool]): This parameter takes effect during inference,
+            indicating whether to enable thinking mode. Default is None, the default value is determined by the
+            template (model) type (True for thinking/hybrid thinking templates, False for non-thinking templates).
+            If enable_thinking is False, a non-thinking prefix is added, for example the Qwen3-8B hybrid thinking
+            model adds the prefix `'<think>\n\n</think>\n\n'`, while Qwen3-8B-Thinking does not add a prefix.
+            If enable_thinking is True, a thinking prefix is added, for example `'<think>\n'`.
+            Note: The priority of this parameter is lower than the response_prefix parameter.
+        preserve_thinking (Optional[bool]): Whether to preserve historical thinking content during inference and
+            training. When set to `True`, thinking content from all rounds is retained. When set to `False`,
+            only the thinking content from the last round is retained (i.e., the content following the last
+            user message). Defaults to `None`.
+            Default behavior: For thinking models (thinking/hybrid-thinking) or when `enable_thinking` is
+            explicitly enabled, this is set to `False` by default during inference and training, retaining
+            only the last round of thinking content. If the `loss_scale` base strategy during training is
+            not `'last_round'` (e.g., `'default'`), it defaults to `True`, and historical thinking content will
+            not be removed.
+        add_non_thinking_prefix (bool): This parameter only takes effect during training, indicating whether to
+            add a non-thinking prefix to data samples whose assistant part does not start with the thinking
+            marker `'<think>'` (typically hybrid thinking models contain a non-thinking prefix).
+            This feature allows swift's built-in datasets to train hybrid thinking models. Default value is True.
+            For example: the non-thinking prefix for the Qwen3-8B hybrid thinking model is
+            `'<think>\n\n</think>\n\n'`, while the non-thinking prefix for Qwen3-8B-Thinking/Instruct is `''`.
+            Note: During training, if the basic strategy of loss_scale is last_round, this modification is only
+            applied to the last round; otherwise, for example 'default' or 'all', this modification is applied to
+            every round of data. If set to False, no non-thinking prefix is added to data samples.
+
+
+    """
+    template: Optional[str] = field(
+        default=None, metadata={'help': f'template choices: {list(TEMPLATE_MAPPING.keys())}'})
+    system: Optional[str] = None  # Override the default_system in the template.
+    max_length: Optional[int] = None
+
+    truncation_strategy: Literal['delete', 'left', 'right', 'split', None] = None
+    max_pixels: Optional[int] = None
+    agent_template: Optional[str] = None
+    norm_bbox: Literal['norm1000', 'none', None] = None
+    use_chat_template: Optional[bool] = None
+    padding_side: Literal['left', 'right'] = 'right'
+    # train
+    padding_free: bool = False
+    loss_scale: str = 'default'
+    sequence_parallel_size: int = 1
+    is_binary_loss_scale: Optional[bool] = None
+    # infer/deploy
+    template_backend: Literal['swift', 'jinja'] = 'swift'
+    # thinking
+    response_prefix: Optional[str] = None
+    enable_thinking: Optional[bool] = None
+    preserve_thinking: Optional[bool] = None
+    add_non_thinking_prefix: bool = True
+    disable_ignore_empty_think: bool = False
+
+    def __post_init__(self):
+        if getattr(self, 'model_meta', None) is not None:
+            self.template_meta = get_template_meta(self.model_info, self.model_meta, template_type=self.template)
+            self.template = self.template_meta.template_type
+        if self.use_chat_template is None:
+            self.use_chat_template = True
+        if self.system is not None:
+            if self.system.endswith('.txt'):
+                assert os.path.isfile(self.system), f'self.system: {self.system}'
+                with open(self.system, 'r', encoding='utf-8') as f:
+                    self.system = f.read()
+            else:
+                self.system = self.system.replace('\\n', '\n')
+        if self.response_prefix is not None:
+            self.response_prefix = self.response_prefix.replace('\\n', '\n')
+        if self.truncation_strategy is None:
+            self.truncation_strategy = 'delete'
+        self._set_loss_scale()
+
+    def _set_loss_scale(self):
+        """For hybrid thinking models, automatically append '+ignore_empty_think' to loss_scale."""
+        if not self.disable_ignore_empty_think and getattr(self, 'template_meta', None) is not None:
+            template_meta = self.template_meta
+            if template_meta.is_thinking and template_meta.non_thinking_prefix:
+                # hybrid thinking model detected
+                if self.loss_scale and 'ignore_empty_think' not in self.loss_scale:
+                    self.loss_scale = self.loss_scale + '+ignore_empty_think'
+
+    def get_template_kwargs(self):
+        truncation_strategy = self.truncation_strategy
+        if truncation_strategy == 'delete':
+            truncation_strategy = 'raise'
+        return {
+            'template_type': self.template,
+            'default_system': self.system,
+            'max_length': self.max_length,
+            'truncation_strategy': truncation_strategy,
+            'max_pixels': self.max_pixels,
+            'agent_template': self.agent_template,
+            'norm_bbox': self.norm_bbox,
+            'use_chat_template': self.use_chat_template,
+            'remove_unused_columns': self.remove_unused_columns,  # from DataArguments
+            'padding_side': self.padding_side,
+            # train
+            'padding_free': self.padding_free,
+            'loss_scale': self.loss_scale,
+            'is_binary_loss_scale': self.is_binary_loss_scale,
+            'sequence_parallel_size': self.sequence_parallel_size,
+            # infer/deploy
+            'template_backend': self.template_backend,
+            # thinking
+            'response_prefix': self.response_prefix,
+            'enable_thinking': self.enable_thinking,
+            'preserve_thinking': self.preserve_thinking,
+            'add_non_thinking_prefix': self.add_non_thinking_prefix,
+        }

swift/arguments/deploy_args.py

@@ -0,0 +1,180 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import os
+from dataclasses import dataclass
+from typing import Literal, Optional
+
+from swift.utils import find_free_port, get_device_count, get_logger, safe_snapshot_download
+from .base_args import BaseArguments
+from .infer_args import InferArguments
+
+logger = get_logger()
+
+
+@dataclass
+class DeployArguments(InferArguments):
+    """Arguments for model deployment.
+
+    This dataclass, which extends InferArguments, is used to define the arguments required for deploying a model.
+
+    Args:
+        host (str): The host address to bind the server to. Defaults to '0.0.0.0'.
+        port (int): The port number to bind the server to. Defaults to 8000.
+        api_key (Optional[str]): The API key for authentication. Defaults to None.
+        ssl_keyfile (Optional[str]): The path to the SSL key file. Defaults to None.
+        ssl_certfile (Optional[str]): The path to the SSL certificate file. Defaults to None.
+        owned_by (str): The owner of the deployment. Defaults to 'swift'.
+        served_model_name (Optional[str]): The name of the model being served. If None, the model's suffix is used by
+            default.
+        verbose (bool): Whether to log detailed request information. Defaults to True.
+            Note: This defaults to False when used in 'swift app' or 'swift eval'.
+        log_interval (int): The interval in seconds for printing tokens/s statistics. Set to -1 to disable. Defaults
+            to 20.
+        log_level (Literal['critical', 'error', 'warning', 'info', 'debug', 'trace']): Log level. Defaults to 'info'.
+        max_logprobs (int): The maximum number of logprobs to return to the client. Defaults to 20.
+        vllm_use_async_engine (Optional[bool]): Whether to use async engine for vLLM.If not set, it defaults to `True`
+            for deployment scenarios.
+    """
+    host: str = '0.0.0.0'
+    port: int = 8000
+    api_key: Optional[str] = None
+    ssl_keyfile: Optional[str] = None
+    ssl_certfile: Optional[str] = None
+
+    owned_by: str = 'swift'
+    served_model_name: Optional[str] = None
+    verbose: bool = True  # Whether to log request_info
+    log_interval: int = 20  # Interval for printing global statistics
+    log_level: Literal['critical', 'error', 'warning', 'info', 'debug', 'trace'] = 'info'
+
+    max_logprobs: int = 20
+    vllm_use_async_engine: Optional[bool] = None
+
+    def __post_init__(self):
+        # default to True for deployment scenarios
+        if self.vllm_use_async_engine is None:
+            self.vllm_use_async_engine = True
+        super().__post_init__()
+        self.port = find_free_port(self.port)
+
+    def _init_adapters(self):
+        if isinstance(self.adapters, str):
+            self.adapters = [self.adapters]
+        self.adapter_mapping = {}
+        adapters = []
+        for i, adapter in enumerate(self.adapters):
+            adapter_path = adapter.split('=')
+            if len(adapter_path) == 1:
+                adapter_path = (None, adapter_path[0])
+            adapter_name, adapter_path = adapter_path
+            adapter_path = safe_snapshot_download(adapter_path, use_hf=self.use_hf, hub_token=self.hub_token)
+            if adapter_name is None:
+                adapters.append(adapter_path)
+            else:
+                self.adapter_mapping[adapter_name] = adapter_path
+        self.adapters = adapters
+
+    def _init_ckpt_dir(self, adapters=None):
+        return super()._init_ckpt_dir(self.adapters + list(self.adapter_mapping.values()))
+
+    def _init_stream(self):
+        return BaseArguments._init_stream(self)
+
+
+@dataclass
+class RolloutArguments(DeployArguments):
+    """Arguments for the Rollout phase in online/reinforcement learning.
+
+    This dataclass inherits from DeployArguments and adds specific parameters for the Rollout process in online
+    learning, such as GRPO.
+
+    Args:
+        multi_turn_scheduler (Optional[str]): The scheduler for multi-turn GRPO training. Pass the name of the
+            corresponding plugin implemented in `swift/rollout/multi_turn.py`. Defaults to None. Refer to the
+            documentation for details.
+        max_turns (Optional[int]): The maximum number of turns in multi-turn GRPO training. If None, no limit is
+            imposed. Defaults to None.
+        vllm_enable_lora (bool): Whether to enable the vLLM Engine to load LoRA adapters. Enabling this can accelerate
+            weight synchronization during LoRA training. Defaults to False. Refer to the documentation for details.
+        vllm_max_lora_rank (int): The LoRA rank parameter for the vLLM Engine. This value must be greater than or
+            equal to the `lora_rank` used for training; setting them as equal is recommended. Defaults to 16.
+    """
+    vllm_use_async_engine: Optional[bool] = None
+    use_gym_env: Optional[bool] = None
+    # only for GRPO rollout with AsyncEngine, see details in swift/rollout/multi_turn
+    multi_turn_scheduler: Optional[str] = None
+    max_turns: Optional[int] = None
+    vllm_enable_lora: bool = False
+    vllm_max_lora_rank: int = 16
+    # GYM env
+    gym_env: Optional[str] = None
+    context_manager: Optional[str] = None
+
+    def __post_init__(self):
+        self._set_default_engine_type()
+        super().__post_init__()
+        self._check_args()
+        self._check_device_count()
+        self._check_vllm_enable_expert_parallel()
+        self._check_deprecated_args()
+        self._set_default_audio_load_backend()
+
+    def _set_default_engine_type(self):
+        if self.vllm_use_async_engine is None:
+            if self.multi_turn_scheduler:
+                self.vllm_use_async_engine = True
+            else:
+                self.vllm_use_async_engine = False
+
+        if self.use_gym_env is None:
+            self.use_gym_env = self.gym_env is not None
+
+    def _check_args(self):
+        if self.vllm_pipeline_parallel_size > 1:
+            raise ValueError('RolloutArguments does not support pipeline parallelism, '
+                             'please set vllm_pipeline_parallel_size to 1.')
+
+        if self.vllm_reasoning_parser is not None:
+            raise ValueError('vllm_reasoning_parser is not supported for Rollout, please unset it.')
+
+        if self.multi_turn_scheduler and not self.vllm_use_async_engine:
+            raise ValueError('please set vllm_use_async_engine to True with multi-turn scheduler.')
+
+    def _check_device_count(self):
+        local_device_count = get_device_count()
+        required_device_count = self.vllm_data_parallel_size * self.vllm_tensor_parallel_size
+
+        if local_device_count < required_device_count:
+            msg = (f'Error: local_device_count ({local_device_count}) must be greater than or equal to '
+                   f'the product of vllm_data_parallel_size ({self.vllm_data_parallel_size}) and '
+                   f'vllm_tensor_parallel_size ({self.vllm_tensor_parallel_size}). '
+                   f'Current required_device_count = {required_device_count}.')
+            raise ValueError(msg)
+
+        if local_device_count > required_device_count:
+            logger.warning_once(
+                f'local_device_count ({local_device_count}) is greater than required_device_count ({required_device_count}). '  # noqa
+                f'Only the first {required_device_count} devices will be utilized for rollout. '
+                f'To fully utilize resources, set vllm_tensor_parallel_size * vllm_data_parallel_size = device_count. '  # noqa
+                f'device_count: {local_device_count}, '
+                f'vllm_tensor_parallel_size: {self.vllm_tensor_parallel_size}, '
+                f'vllm_data_parallel_size: {self.vllm_data_parallel_size}, '
+                f'required_device_count: {required_device_count}.')
+
+    def _check_vllm_enable_expert_parallel(self):
+        if self.vllm_enable_expert_parallel and not self.vllm_use_async_engine:
+            self.vllm_use_async_engine = True
+            logger.warning('vllm_enable_expert_parallel is only supported with vllm_use_async_engine, '
+                           'set vllm_use_async_engine to True.')
+
+    def _check_deprecated_args(self):
+        if self.context_manager is not None:
+            raise ValueError('The "context_manager" argument has been removed. '
+                             'If you need to dynamically modify the conversation history between rollout turns '
+                             '(e.g. history compression, prompt injection), implement that logic in a custom '
+                             '`MultiTurnScheduler` subclass by overriding `step` / `run`, '
+                             'and pass it via `--multi_turn_scheduler your_scheduler_name`.')
+
+    def _set_default_audio_load_backend(self):
+        # Rollout uses GRPOVllmEngine (vLLM-only); align audio decode with vLLM multimodal loader.
+        if os.getenv('SWIFT_AUDIO_LOAD_BACKEND') is None:
+            os.environ['SWIFT_AUDIO_LOAD_BACKEND'] = 'soundfile_pyav'

swift/arguments/eval_args.py

@@ -0,0 +1,130 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import datetime as dt
+import os
+from dataclasses import dataclass, field
+from typing import List, Literal, Optional, Union
+
+from swift.model import get_matched_model_meta
+from swift.utils import get_logger, json_parse_to_dict, to_abspath
+from .deploy_args import DeployArguments
+
+logger = get_logger()
+
+
+@dataclass
+class EvalArguments(DeployArguments):
+    """A dataclass that extends DeployArguments to define model evaluation arguments.
+
+    These arguments control the evaluation process, including the choice of backend, datasets, generation parameters,
+    and other configurations.
+
+    Args:
+        eval_dataset (List[str]): List of evaluation datasets. Please refer to the evaluation documentation for
+            available options. Defaults to [].
+        eval_limit (Optional[int]): The number of samples to take from each evaluation dataset. If None, all samples
+            are used. Defaults to None.
+        eval_dataset_args (Optional[Union[Dict, str]]): Evaluation dataset parameters, in JSON format, can be set for
+            multiple datasets. Defaults to None.
+        eval_generation_config (Optional[Union[Dict, str]]): The model's inference configuration for evaluation,
+            provided as a JSON string (e.g., '{"max_new_tokens": 512}'). Defaults to None.
+        eval_output_dir (str): The directory to store evaluation results. Defaults to 'eval_output'.
+        eval_backend (str): The evaluation backend. Can be 'Native', 'OpenCompass', or 'VLMEvalKit'. Defaults to
+            'Native'.
+        local_dataset (bool): Whether to automatically download extra datasets required for certain evaluations
+            (e.g., CMB). If True, a 'data' folder will be created in the current directory for the datasets. This
+            download occurs only once, and subsequent runs will use the cache. Defaults to False.
+            Note: By default, evaluation uses datasets from `~/.cache/opencompass`. When this is set to True, the
+            `data` folder in the current directory is used instead.
+        temperature (float): The temperature for sampling, which overrides the default generation config. Defaults
+            to 0.0.
+        verbose (bool): Whether to output verbose information during the evaluation process. Defaults to False.
+        eval_num_proc (int): The maximum number of concurrent clients for evaluation. Defaults to 16.
+        extra_eval_args (Optional[Union[Dict, str]]): Additional evaluation arguments, provided as a JSON string.
+            These are only effective when using the 'Native' backend. Refer to the documentation for more details on
+            available arguments. Defaults to {}.
+        eval_url (Optional[str]): The URL for the evaluation service (e.g., 'http://localhost:8000/v1'). If not
+            specified, evaluation runs on the locally deployed model. See documentation for more examples. Defaults
+            to None.
+    """
+    eval_dataset: List[str] = field(default_factory=list)
+    eval_limit: Optional[int] = None
+    eval_dataset_args: Optional[Union[dict, str]] = None
+    eval_generation_config: Optional[Union[dict, str]] = None
+    eval_output_dir: str = 'eval_output'
+    eval_backend: Literal['Native', 'OpenCompass', 'VLMEvalKit'] = 'Native'
+    local_dataset: bool = False
+
+    temperature: Optional[float] = 0.
+    verbose: bool = False
+    eval_num_proc: int = 16
+    extra_eval_args: Optional[Union[dict, str]] = field(default_factory=dict)
+    # If eval_url is set, ms-swift will not perform deployment operations and
+    # will directly use the URL for evaluation.
+    eval_url: Optional[str] = None
+
+    def __post_init__(self):
+        super().__post_init__()
+        self._init_eval_url()
+        self._init_eval_dataset()
+        self.eval_dataset_args = json_parse_to_dict(self.eval_dataset_args)
+        self.eval_generation_config = json_parse_to_dict(self.eval_generation_config)
+        self.extra_eval_args = json_parse_to_dict(self.extra_eval_args)
+        self.eval_output_dir = to_abspath(self.eval_output_dir)
+        logger.info(f'eval_output_dir: {self.eval_output_dir}')
+
+    def _init_eval_url(self):
+        # [compat]
+        if self.eval_url and 'chat/completions' in self.eval_url:
+            self.eval_url = self.eval_url.split('/chat/completions', 1)[0]
+
+    @staticmethod
+    def list_eval_dataset(eval_backend=None):
+        from evalscope.api.registry import BENCHMARK_REGISTRY
+        from evalscope.backend.opencompass import OpenCompassBackendManager
+        from evalscope.constants import EvalBackend
+        res = {
+            EvalBackend.NATIVE: list(sorted(BENCHMARK_REGISTRY.keys())),
+            EvalBackend.OPEN_COMPASS: sorted(OpenCompassBackendManager.list_datasets()),
+        }
+        try:
+            from evalscope.backend.vlm_eval_kit import VLMEvalKitBackendManager
+            vlm_datasets = VLMEvalKitBackendManager.list_supported_datasets()
+            res[EvalBackend.VLM_EVAL_KIT] = sorted(vlm_datasets)
+        except ImportError:
+            # fix cv2 import error
+            if eval_backend == 'VLMEvalKit':
+                raise
+        return res
+
+    def _init_eval_dataset(self):
+        if isinstance(self.eval_dataset, str):
+            self.eval_dataset = [self.eval_dataset]
+
+        all_eval_dataset = self.list_eval_dataset(self.eval_backend)
+        dataset_mapping = {dataset.lower(): dataset for dataset in all_eval_dataset[self.eval_backend]}
+        valid_dataset = []
+        for dataset in self.eval_dataset:
+            if dataset.lower() not in dataset_mapping:
+                raise ValueError(
+                    f'eval_dataset: {dataset} is not supported.\n'
+                    f'eval_backend: {self.eval_backend} supported datasets: {all_eval_dataset[self.eval_backend]}')
+            valid_dataset.append(dataset_mapping[dataset.lower()])
+        self.eval_dataset = valid_dataset
+
+        logger.info(f'eval_backend: {self.eval_backend}')
+        logger.info(f'eval_dataset: {self.eval_dataset}')
+
+    def _init_result_path(self, folder_name: str) -> None:
+        self.time = dt.datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')
+        result_dir = self.ckpt_dir or f'result/{self.model_suffix}'
+        os.makedirs(result_dir, exist_ok=True)
+        self.result_jsonl = to_abspath(os.path.join(result_dir, 'eval_result.jsonl'))
+        if not self.eval_url:
+            super()._init_result_path('eval_result')
+
+    def _init_torch_dtype(self) -> None:
+        if self.eval_url:
+            self.model_meta = get_matched_model_meta(self.model)
+            self.model_info = None
+            return
+        super()._init_torch_dtype()

swift/arguments/export_args.py

@@ -0,0 +1,150 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import os
+import torch
+import torch.distributed as dist
+from dataclasses import dataclass
+from typing import Literal, Optional
+
+from swift.utils import HfConfigFactory, get_logger, init_process_group, set_default_ddp_config, to_abspath
+from .base_args import BaseArguments
+from .merge_args import MergeArguments
+
+logger = get_logger()
+
+
+@dataclass
+class ExportArguments(MergeArguments, BaseArguments):
+    """ExportArguments is a dataclass that inherits from BaseArguments and MergeArguments.
+
+    Args:
+        output_dir (Optional[str]): Directory to save the exported results. Defaults to None, which automatically sets
+            a path with an appropriate suffix.
+        quant_method (Optional[str]): The quantization method. Can be 'awq', 'gptq', 'bnb', 'fp8', or 'gptq_v2'.
+            Defaults to None. See examples for more details.
+        quant_n_samples (int): Number of samples for GPTQ/AWQ calibration. Defaults to 256.
+        quant_batch_size (int): The batch size for quantization. Defaults to 1.
+        group_size (int): The group size for quantization. Defaults to 128.
+        to_cached_dataset (bool): Whether to tokenize and export the dataset in advance as a cached dataset. Defaults
+            to False. Note: You can specify the validation set content through
+            `--split_dataset_ratio` or `--val_dataset`.
+        to_ollama (bool): Whether to generate the `Modelfile` required by Ollama. Defaults to False.
+        to_mcore (bool): Whether to convert Hugging Face format weights to Megatron-Core format. Defaults to False.
+        to_hf (bool): Whether to convert Megatron-Core format weights to Hugging Face format. Defaults to False.
+        mcore_model (Optional[str]): The path to the Megatron-Core format model. Defaults to None.
+        mcore_adapter (Optional[str]): A list of adapter paths for the Megatron-Core format model. Defaults to [].
+        thread_count (Optional[int]): The number of model shards when `to_mcore` is True. Defaults to None, which
+            automatically sets the number based on the model size to keep the largest shard under 10GB.
+        test_convert_precision (bool): Whether to test the precision error of weight conversion between Hugging Face
+            and Megatron-Core formats. Defaults to False.
+        test_convert_dtype (str): The dtype to use for the conversion precision test. Defaults to 'float32'.
+        push_to_hub (bool): Whether to push the output to the Model Hub. Defaults to False. See examples for more
+            details.
+        hub_model_id (Optional[str]): The model ID for pushing to the Hub (e.g., 'user_name/repo_name' or 'repo_name').
+            Defaults to None.
+        hub_private_repo (bool): Whether the Hub repository is private. Defaults to False.
+        commit_message (str): The commit message for pushing to the Hub. Defaults to 'update files'.
+        to_peft_format (bool): Whether to export in PEFT format. This argument is for compatibility and currently has
+            no effect. Defaults to False.
+        exist_ok (bool): If the output_dir exists, do not raise an exception and overwrite its contents. Defaults to
+            False.
+    """
+    output_dir: Optional[str] = None
+
+    # awq/gptq
+    quant_method: Literal['awq', 'gptq', 'bnb', 'fp8', 'gptq_v2'] = None
+    quant_n_samples: int = 256
+    quant_batch_size: int = 1
+    group_size: int = 128
+
+    # cached_dataset
+    to_cached_dataset: bool = False
+    template_mode: Literal['train', 'rlhf', 'kto'] = 'train'
+
+    # ollama
+    to_ollama: bool = False
+
+    # megatron
+    to_mcore: bool = False
+    to_hf: bool = False
+    mcore_model: Optional[str] = None
+    mcore_adapter: Optional[str] = None
+    thread_count: Optional[int] = None
+    test_convert_precision: bool = False
+    test_convert_dtype: str = 'float32'
+
+    # push to ms hub
+    push_to_hub: bool = False
+    # 'user_name/repo_name' or 'repo_name'
+    hub_model_id: Optional[str] = None
+    hub_private_repo: bool = False
+    commit_message: str = 'update files'
+    # compat
+    to_peft_format: bool = False
+    exist_ok: bool = False
+
+    def load_args_from_ckpt(self) -> None:
+        if self.to_cached_dataset:
+            return
+        super().load_args_from_ckpt()
+
+    def _init_output_dir(self):
+        if self.output_dir is None:
+            ckpt_dir = self.ckpt_dir or f'./{self.model_suffix}'
+            ckpt_dir, ckpt_name = os.path.split(ckpt_dir)
+            if self.to_peft_format:
+                suffix = 'peft'
+            elif self.quant_method:
+                suffix = f'{self.quant_method}'
+                if self.quant_bits is not None:
+                    suffix += f'-int{self.quant_bits}'
+            elif self.to_ollama:
+                suffix = 'ollama'
+            elif self.merge_lora:
+                suffix = 'merged'
+            elif self.to_mcore:
+                suffix = 'mcore'
+            elif self.to_hf:
+                suffix = 'hf'
+            elif self.to_cached_dataset:
+                suffix = 'cached_dataset'
+            else:
+                return
+
+            self.output_dir = os.path.join(ckpt_dir, f'{ckpt_name}-{suffix}')
+
+        self.output_dir = to_abspath(self.output_dir)
+        if not self.exist_ok and os.path.exists(self.output_dir):
+            raise FileExistsError(f'args.output_dir: `{self.output_dir}` already exists.')
+        logger.info(f'args.output_dir: `{self.output_dir}`')
+
+    def __post_init__(self):
+        if self.quant_batch_size == -1:
+            self.quant_batch_size = None
+        if self.quant_bits and self.quant_method is None:
+            raise ValueError('Please specify the quantization method using `--quant_method awq/gptq/bnb`.')
+        if self.quant_method and self.quant_bits is None and self.quant_method != 'fp8':
+            raise ValueError('Please specify `--quant_bits`.')
+        if self.quant_method in {'gptq', 'awq'} and self.torch_dtype is None:
+            self.torch_dtype = torch.float16
+        if self.to_mcore or self.to_hf:
+            if self.merge_lora:
+                self.merge_lora = False
+                logger.warning('`swift export --to_mcore/to_hf` does not support the `--merge_lora` parameter. '
+                               'To export LoRA delta weights, please use `megatron export`')
+
+            self.mcore_model = to_abspath(self.mcore_model, check_path_exist=True)
+            if not dist.is_initialized():
+                set_default_ddp_config()
+                init_process_group(backend=self.ddp_backend, timeout=self.ddp_timeout)
+
+        BaseArguments.__post_init__(self)
+        self._init_output_dir()
+        self.test_convert_dtype = HfConfigFactory.to_torch_dtype(self.test_convert_dtype)
+        if self.quant_method in {'gptq', 'awq'} and len(self.dataset) == 0:
+            raise ValueError(f'self.dataset: {self.dataset}, Please input the quant dataset.')
+        if self.to_cached_dataset:
+            self.lazy_tokenize = False
+            if self.packing:
+                raise ValueError('Packing will be handled during training; here we only perform tokenization '
+                                 'in advance, so you do not need to set up packing separately.')
+            assert not self.streaming, 'not supported'

swift/arguments/infer_args.py

@@ -0,0 +1,237 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import datetime as dt
+import os
+import torch.distributed as dist
+from dataclasses import dataclass
+from typing import Literal, Optional
+
+from swift.rlhf_trainers import VllmArguments
+from swift.utils import get_logger, init_process_group, is_dist, to_abspath
+from .base_args import BaseArguments
+from .merge_args import MergeArguments
+
+logger = get_logger()
+
+
+@dataclass
+class LmdeployArguments:
+    """Holds the configuration arguments for lmdeploy.
+
+    Args:
+        lmdeploy_tp (int): The tensor parallelism size. Defaults to 1.
+        lmdeploy_session_len (Optional[int]): The maximum session length. Defaults to None.
+        lmdeploy_cache_max_entry_count (float): The percentage of GPU memory to be used by the K/V cache. Defaults
+            to 0.8.
+        lmdeploy_quant_policy (int): The quantization policy for the K/V cache. Set to 4 or 8 for 4-bit or 8-bit
+            quantization respectively. Defaults to 0, which means no quantization.
+        lmdeploy_vision_batch_size (int): The `max_batch_size` parameter to be passed to `VisionConfig`. Defaults to 1.
+    """
+
+    # lmdeploy
+    lmdeploy_tp: int = 1
+    lmdeploy_session_len: Optional[int] = None
+    lmdeploy_cache_max_entry_count: float = 0.8
+    lmdeploy_quant_policy: int = 0  # e.g. 4, 8
+    lmdeploy_vision_batch_size: int = 1  # max_batch_size in VisionConfig
+
+    def get_lmdeploy_engine_kwargs(self):
+        kwargs = {
+            'tp': self.lmdeploy_tp,
+            'session_len': self.lmdeploy_session_len,
+            'cache_max_entry_count': self.lmdeploy_cache_max_entry_count,
+            'quant_policy': self.lmdeploy_quant_policy,
+            'vision_batch_size': self.lmdeploy_vision_batch_size
+        }
+        if dist.is_initialized():
+            kwargs.update({'devices': [dist.get_rank()]})
+        return kwargs
+
+
+@dataclass
+class SglangArguments:
+    """Arguments for configuring the SGLang backend.
+
+    Args:
+        sglang_tp_size (int): The number of tensor parallel workers. Defaults to 1.
+        sglang_pp_size (int): The number of pipeline parallel workers. Defaults to 1.
+        sglang_dp_size (int): The number of data parallel workers. Defaults to 1.
+        sglang_ep_size (int): The number of expert parallel workers. Defaults to 1.
+        sglang_enable_ep_moe (bool): Whether to enable expert parallelism for MoE.
+            Note: This argument has been removed in recent versions of SGLang. Defaults to False.
+        sglang_mem_fraction_static (Optional[float]): The fraction of GPU memory for the static allocation of model
+            weights and the KV cache memory pool. Try lowering this value if you encounter GPU out-of-memory errors.
+            Defaults to None.
+        sglang_context_length (Optional[int]): The maximum context length for the model. If None, the value from the
+            model's `config.json` will be used. Defaults to None.
+        sglang_disable_cuda_graph (bool): Disable CUDA graph for inference. Defaults to False.
+        sglang_quantization (Optional[str]): The quantization method to use. Defaults to None.
+        sglang_kv_cache_dtype (str): The data type for K/V cache storage. 'auto' will use the model's data type.
+            'fp8_e5m2' and 'fp8_e4m3' are available for CUDA 11.8 and later. Defaults to 'auto'.
+        sglang_enable_dp_attention (bool): Enables data parallelism for the attention mechanism and tensor parallelism
+            for the feed-forward network (FFN). The data parallel size (dp_size) must equal the tensor parallel size
+            (tp_size). Currently supported for DeepSeek-V2/3 and Qwen2/3 MoE models. Defaults to False.
+        sglang_disable_custom_all_reduce (bool): Disable the custom all-reduce kernel and fall back to NCCL. Enabled by
+            default (True) for stability. Defaults to True.
+        sglang_speculative_algorithm (Optional[str]): The speculative decoding algorithm. Options include "EAGLE",
+            "EAGLE3", "NEXTN", "STANDALONE", "NGRAM". Defaults to None.
+        sglang_speculative_num_steps (Optional[int]): The number of steps to sample from the draft model during
+            speculative decoding. Defaults to None.
+        sglang_speculative_eagle_topk (Optional[int]): The number of tokens to sample from the draft model at each step
+            for the EAGLE2 algorithm. Defaults to None.
+        sglang_speculative_num_draft_tokens (Optional[int]): The number of tokens to sample from the draft model during
+            speculative decoding. Defaults to None.
+    """
+    sglang_tp_size: int = 1
+    sglang_pp_size: int = 1
+    sglang_dp_size: int = 1
+    sglang_ep_size: int = 1
+    sglang_enable_ep_moe: bool = False
+    sglang_mem_fraction_static: Optional[float] = None
+    sglang_context_length: Optional[int] = None
+    sglang_disable_cuda_graph: bool = False
+    sglang_quantization: Optional[str] = None
+    sglang_kv_cache_dtype: str = 'auto'
+    sglang_enable_dp_attention: bool = False
+    sglang_disable_custom_all_reduce: bool = True
+    # speculative decoding
+    # e.g. EAGLE, EAGLE3, NEXTN
+    sglang_speculative_algorithm: Optional[str] = None
+    sglang_speculative_num_steps: Optional[int] = None
+    sglang_speculative_eagle_topk: Optional[int] = None
+    sglang_speculative_num_draft_tokens: Optional[int] = None
+
+    def get_sglang_engine_kwargs(self):
+        kwargs = {
+            'tp_size': self.sglang_tp_size,
+            'pp_size': self.sglang_pp_size,
+            'dp_size': self.sglang_dp_size,
+            'ep_size': self.sglang_ep_size,
+            'enable_ep_moe': self.sglang_enable_ep_moe,
+            'mem_fraction_static': self.sglang_mem_fraction_static,
+            'context_length': self.sglang_context_length,
+            'disable_cuda_graph': self.sglang_disable_cuda_graph,
+            'quantization': self.sglang_quantization,
+            'kv_cache_dtype': self.sglang_kv_cache_dtype,
+            'enable_dp_attention': self.sglang_enable_dp_attention,
+            'disable_custom_all_reduce': self.sglang_disable_custom_all_reduce,
+            'speculative_algorithm': self.sglang_speculative_algorithm,
+            'speculative_num_steps': self.sglang_speculative_num_steps,
+            'speculative_eagle_topk': self.sglang_speculative_eagle_topk,
+            'speculative_num_draft_tokens': self.sglang_speculative_num_draft_tokens,
+        }
+        if self.task_type == 'embedding':
+            kwargs['task_type'] = 'embedding'
+        return kwargs
+
+
+@dataclass
+class InferArguments(MergeArguments, LmdeployArguments, SglangArguments, VllmArguments, BaseArguments):
+    """Arguments for model inference.
+
+    A dataclass that extends BaseArguments, MergeArguments, VllmArguments, and LmdeployArguments to define all
+    arguments required for model inference.
+
+    Args:
+        infer_backend (Literal['transformers', 'vllm', 'sglang', 'lmdeploy']): The inference acceleration
+            backend to use. Defaults to 'transformers'.
+        result_path (Optional[str]): The path to store inference results in JSONL format. If the file already exists,
+            new results will be appended. If None, results are saved in the checkpoint directory (if available) or
+            './result'. The final path will be printed to the console. Defaults to None.
+        write_batch_size (int): The batch size for writing results to `result_path`. A value of -1 means no limit.
+            Defaults to 1000.
+        metric (Optional[str]): The metric to use for evaluating inference results. Supported values are 'acc' and
+            'rouge'. If None, no evaluation is performed. Defaults to None.
+        max_batch_size (int): The maximum batch size for inference, effective only when `infer_backend` is
+            'transformers'. A value of -1 means no limit. Defaults to 1.
+        val_dataset_sample (Optional[int]): The number of samples to use from the inference dataset. If None, the
+            entire dataset is used. Defaults to None.
+        reranker_use_activation (bool): Whether to apply a sigmoid activation to the scores during reranker inference.
+            Defaults to True.
+    """
+    # `pt` is used for swift3.x shell script compatibility.
+    infer_backend: Literal['vllm', 'transformers', 'sglang', 'lmdeploy', 'pt'] = 'transformers'
+
+    result_path: Optional[str] = None
+    write_batch_size: int = 1000
+    metric: Literal['acc', 'rouge'] = None
+    # for transformers engine
+    max_batch_size: int = 1
+
+    # only for inference
+    val_dataset_sample: Optional[int] = None
+
+    # for reranker
+    reranker_use_activation: bool = True
+
+    def _get_result_path(self, folder_name: str) -> str:
+        result_dir = self.ckpt_dir or f'result/{self.model_suffix}'
+        os.makedirs(result_dir, exist_ok=True)
+        result_dir = to_abspath(os.path.join(result_dir, folder_name))
+        os.makedirs(result_dir, exist_ok=True)
+        time = dt.datetime.now().strftime('%Y%m%d-%H%M%S')
+        return os.path.join(result_dir, f'{time}.jsonl')
+
+    def _init_result_path(self, folder_name: str) -> None:
+        if self.result_path is not None:
+            self.result_path = to_abspath(self.result_path)
+            return
+        # By default, a result_path file is automatically created
+        # when a validation or evaluation dataset is present.
+        if self._val_dataset_exists or getattr(self, 'eval_dataset', None):
+            self.result_path = self._get_result_path(folder_name)
+            logger.info(f'args.result_path: {self.result_path}')
+
+    def _init_stream(self):
+        self.eval_human = not self._val_dataset_exists
+        logger.info(f'Setting args.eval_human: {self.eval_human}')
+        if self.stream is None:
+            self.stream = self.eval_human
+        if self.stream and self.num_beams != 1:
+            self.stream = False
+            logger.info('Setting args.stream: False')
+
+    def _init_ddp(self):
+        if not is_dist():
+            return
+        eval_human = getattr(self, 'eval_human', False)
+        assert not eval_human and not self.stream, (
+            'In DDP scenarios, interactive interfaces and streaming output are not supported.'
+            f'args.eval_human: {eval_human}, args.stream: {self.stream}')
+        self._init_device()
+        init_process_group(backend=self.ddp_backend, timeout=self.ddp_timeout)
+
+    def __post_init__(self) -> None:
+        if self.infer_backend == 'pt':
+            self.infer_backend = 'transformers'  # compat swift3.x
+            logger.warning('args.infer_backend: `pt` is deprecated, please use args.infer_backend: `transformers`.')
+        BaseArguments.__post_init__(self)
+        VllmArguments.__post_init__(self)
+        self._init_vllm_async_engine()
+        # Default to False for swift infer (non-encode tasks)
+        if self.vllm_use_async_engine is None:
+            self.vllm_use_async_engine = False
+        self._init_result_path('infer_result')
+        self._init_ddp()
+
+    def _init_vllm_async_engine(self):
+        """Initialize vllm_use_async_engine based on task_type.
+
+        Encode tasks (embedding, seq_cls, reranker, generative_reranker) require
+        async engine because vLLM's synchronous LLMEngine does not have the `encode` method.
+
+        Note: This method only handles encode tasks. For non-encode tasks, the default value
+        should be set by subclasses (DeployArguments sets True, RolloutArguments uses
+        _set_default_engine_type, InferArguments defaults to False).
+        """
+        # Task types that require vLLM's encode() method, which is only available in AsyncLLMEngine
+        encode_task_types = ('embedding', 'seq_cls', 'reranker', 'generative_reranker')
+        is_vllm_encode_task = self.infer_backend == 'vllm' and self.task_type in encode_task_types
+
+        if is_vllm_encode_task:
+            if self.vllm_use_async_engine is None:
+                self.vllm_use_async_engine = True
+            elif not self.vllm_use_async_engine:
+                raise ValueError(
+                    f'task_type={self.task_type} requires vllm_use_async_engine=True. '
+                    f'The synchronous vLLM LLMEngine does not support the `encode` method for encode tasks. '
+                    f'Please set --vllm_use_async_engine true or remove the explicit false setting.')

swift/arguments/merge_args.py

@@ -0,0 +1,23 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from dataclasses import dataclass
+
+from swift.utils import get_logger
+
+logger = get_logger()
+
+
+@dataclass
+class MergeArguments:
+    """A dataclass that holds configuration for merging models.
+
+    This dataclass stores all the arguments needed to configure the model merging process.
+
+    Args:
+        merge_lora (bool): Whether to merge LoRA adapters. This parameter supports `lora`, `llamapro`, and `longlora`.
+            Defaults to False.
+        safe_serialization (bool): Whether to use safetensors for serialization. Defaults to True.
+        max_shard_size (str): The maximum size of a single saved shard file. Defaults to '5GB'.
+    """
+    merge_lora: bool = False
+    safe_serialization: bool = True
+    max_shard_size: str = '5GB'

swift/arguments/pretrain_args.py

@@ -0,0 +1,10 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from dataclasses import dataclass
+
+from .sft_args import SftArguments
+
+
+@dataclass
+class PretrainArguments(SftArguments):
+    use_chat_template: bool = False
+    loss_scale: str = 'all'

swift/arguments/rlhf_args.py

@@ -0,0 +1,675 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import os
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Literal, Optional
+
+from swift.model import MODEL_MAPPING
+from swift.rlhf_trainers import GRPOArgumentsMixin
+from swift.template import TEMPLATE_MAPPING
+from swift.utils import get_current_device, get_logger, is_master, is_mp, json_parse_to_dict, set_default_ddp_config
+from .sft_args import SftArguments
+
+logger = get_logger()
+rlhf_support_vllm_types = ['grpo', 'gkd']
+
+
+@dataclass
+class RewardModelArguments:
+    """Arguments pertaining to the reward model.
+
+    Args:
+        reward_model (Optional[List[str]]): The model ID or a local path to the reward model. Same as the `model`
+            argument. Defaults to None.
+        reward_adapters (List[str]): The path(s) to LoRA adapter weights to be loaded for the reward model. Useful for
+            using LoRA weights from SFT as the reward model. Defaults to an empty list (`[]`).
+        reward_model_type (Optional[List[str]]): The model type of the reward model. Same as the `model_type` argument.
+            If not specified, it's often inferred. Defaults to None.
+        reward_model_revision (Optional[List[str]]): The specific model version to use for the reward model. Same as
+            the `model_revision` argument. Defaults to None.
+        reward_template (Optional[List[str]]): The template to use for the reward model. Defaults to None.
+    """
+    reward_model: Optional[List[str]] = None
+    reward_adapters: List[str] = field(default_factory=list)
+    reward_model_type: Optional[List[str]] = field(
+        default=None, metadata={'help': f'model_type choices: {list(MODEL_MAPPING.keys())}'})
+    reward_model_revision: Optional[List[str]] = None
+    reward_template: Optional[List[str]] = field(
+        default=None, metadata={'help': f'template choices: {list(TEMPLATE_MAPPING.keys())}'})
+
+
+@dataclass
+class TeacherModelArguments:
+    """Arguments for configuring the teacher model.
+
+    Args:
+        teacher_model (Optional[str]): The model ID or a local path to the teacher model. Analogous to the main
+            `model` argument. For GKD, there are three modes:
+            - Not set (None): Self-distillation with dynamic teacher (teacher = current student weights).
+            - Same as `model` with LoRA training: Self-distillation with fixed teacher. Automatically optimized
+              to use `disable_adapter()` to get base model logits without loading an extra model.
+            - Different from `model`: Standard GKD with an independent frozen teacher model.
+            Defaults to None.
+        teacher_adapters (List[str]): A list of paths to LoRA weights. These weights, often produced by SFT, are loaded
+            to form the teacher model. Defaults to an empty list (`[]`).
+        teacher_model_type (Optional[str]): The model type of the teacher model. If not specified, it's often inferred.
+            Analogous to the main `model_type` argument. Defaults to None.
+        teacher_model_revision (Optional[str]): The specific model version of the teacher model to use. Analogous to
+            the main `model_revision` argument. Defaults to None.
+        teacher_deepspeed (Optional[str]): The teacher model's deepspeed configuration. This can be a JSON file path or
+            one of the following values: 'zero0', 'zero1', 'zero2', 'zero3', 'zero2_offload', 'zero3_offload'. If not
+            provided, it defaults to using the same DeepSpeed configuration as the main training model. Analogous to
+            the main `deepspeed` argument.
+        teacher_model_server (Optional[str]): The URL of the teacher model server (e.g., 'http://localhost:8000').
+            When set, the teacher logprobs will be fetched from the external API service (e.g., swift deploy, vLLM)
+            instead of loading a local teacher model. This enables using larger teacher models or services hosted
+            remotely. When this is set, `teacher_model` is not required. Defaults to None.
+        offload_teacher_model (bool): Whether to offload the teacher model to CPU memory to save VRAM during GKD
+            or OPD-RL training. When enabled, the teacher model is loaded to GPU only during forward pass and
+            offloaded back to CPU afterwards. Defaults to False.
+    """
+    teacher_model: Optional[str] = None
+    teacher_adapters: List[str] = field(default_factory=list)
+    teacher_model_type: Optional[str] = field(
+        default=None, metadata={'help': f'model_type choices: {list(MODEL_MAPPING.keys())}'})
+    teacher_model_revision: Optional[str] = None
+    teacher_deepspeed: Optional[str] = field(
+        default=None,
+        metadata={
+            'help':
+            'DeepSpeed configuration for teacher model. '
+            'Can be a path to a json file or one of: zero0, zero1, zero2, zero3, zero2_offload, zero3_offload'
+        })
+    teacher_model_server: Optional[str] = field(
+        default=None,
+        metadata={
+            'help':
+            'URL of the teacher model server (e.g., http://localhost:8000). '
+            'When set, teacher logprobs are fetched via API instead of loading a local model. '
+            'Supports multi-teacher via JSON list of {url, tags}.'
+        })
+    offload_teacher_model: bool = False
+
+
+@dataclass
+class PPOArguments:
+    """Arguments for configuring the PPO training.
+
+    Args:
+        num_ppo_epochs (int): Number of epochs to train. Defaults to 4.
+        whiten_rewards (bool): Whether to whiten the rewards. Defaults to False.
+        kl_coef (float): KL coefficient. Defaults to 0.05.
+        cliprange (float): Clip range. Defaults to 0.2.
+        vf_coef (float): Value function coefficient. Defaults to 0.1.
+        cliprange_value (float): Clip range for the value function. Defaults to 0.2.
+        gamma (float): Discount factor. Defaults to 1.0.
+        lam (float): Lambda value for GAE. Defaults to 0.95.
+        num_mini_batches (int): Defaults to 1.
+        local_rollout_forward_batch_size (int): Defaults to 64.
+        num_sample_generations (int): Number of generations. Defaults to 10.
+        response_length (Optional[int]): (Deprecated) Compatibility parameter. Use `max_completion_length` instead.
+            Defaults to None.
+        missing_eos_penalty (Optional[float]): Defaults to None.
+    """
+    num_ppo_epochs: int = 4
+    whiten_rewards: bool = False
+    kl_coef: float = 0.05
+    cliprange: float = 0.2
+    vf_coef: float = 0.1
+    cliprange_value: float = 0.2
+    gamma: float = 1.0
+    lam: float = 0.95
+
+    num_mini_batches: int = 1
+    local_rollout_forward_batch_size: int = 64
+    num_sample_generations: int = 10
+    response_length: Optional[int] = None  # compat. use max_completion_length instead
+    missing_eos_penalty: Optional[float] = None
+
+
+@dataclass
+class GRPOArguments(GRPOArgumentsMixin):
+    """A dataclass for configuring GRPO training.
+
+    These arguments control the hyperparameters specific to the GRPO algorithm.
+
+    Args:
+        num_generations  (int): The number of completions to generate for each prompt. This corresponds to the G value
+            in the GRPO paper. The total generation batch size (e.g., `generation_batch_size` or `steps_per_generation
+            * per_device_batch_size * num_processes`) must be divisible by this number. Defaults to 8.
+        num_generations_eval (Optional[int]): Number of generations to sample during evaluation. This allows
+            using fewer generations during evaluation to save computation. If `None`, uses the value of
+            `num_generations`. Defaults to None.
+        reward_funcs (List[str]): A list of reward function names to use for the GRPO algorithm. Available built-in
+            options include 'accuracy', 'format', 'cosine', 'repetition', and 'soft_overlong'
+            (see swift/rewards/orm.py). Custom reward functions can also be defined. Defaults to an empty list.
+        reward_weights (List[float]): A list of weights for each reward source. The length must match the total number
+            of reward functions (from `reward_funcs`) plus any external reward models. If `None`, all rewards are
+            weighted equally with a value of 1.0. Note: If an external `--reward_model` is used, it is treated as the
+            last reward source in the sequence. Defaults to None.
+        log_completions (bool): Whether to log the model's generated completions during training. This is designed to
+            be used with an experiment tracker like WandB or SwanLab (`--report_to wandb`/`swanlab`). If enabled
+            without a tracker, completions are saved to `completions.jsonl` in the checkpoint directory. Defaults to
+            False.
+        num_iterations (int): The number of update steps to perform for each data sample. This corresponds to the K
+            value in the GRPO paper. Defaults to 1.
+        truncation_strategy (Literal['delete', 'left', 'right', 'split', None]): The strategy for handling input
+            sequences that exceed `max_length`. Supported options: 'delete' to discard the sample, 'left' to truncate
+            from the beginning, 'right' to truncate from the end. Defaults to None, and then sets to 'left' in the
+            `_init_grpo` function.
+            Note that for multimodal models, left pruning may prune multimodal tokens, causing shape mismatch errors
+            in the forward feed. Using the `delete` method will resample other data from the original dataset to
+            supplement excessively long data and examples with encoding failures.
+    """
+    num_generations: int = 8  # G in the GRPO paper
+    reward_funcs: List[str] = field(default_factory=list)
+    reward_weights: List[float] = None
+    log_completions: bool = False
+
+    # multi step
+    num_iterations: int = 1
+
+    truncation_strategy: Literal['delete', 'left', 'right', 'split', None] = None
+
+
+@dataclass
+class RLHFArguments(TeacherModelArguments, GRPOArguments, PPOArguments, RewardModelArguments, SftArguments):
+    """A dataclass holding arguments for Reinforcement Learning from Human Feedback.
+
+    Args:
+        rlhf_type (str): The type of human alignment algorithm to use. Supports 'dpo', 'orpo', 'simpo', 'kto', 'cpo',
+            'rm', 'ppo', 'grpo', and 'gkd'. Defaults to 'dpo'.
+        ref_model (Optional[str]): The model path for the reference model. Required when using 'dpo', 'kto', 'ppo',
+            or 'grpo' with full-parameter training. Defaults to None, which will set it to the value of the `--model`
+            argument.
+        ref_adapters (List[str]): LoRA adapters for the reference model. If you are using LoRA weights from SFT for
+            DPO/KTO/GRPO, set both `--adapters` and `--ref_adapters` to the SFT checkpoint path. When resuming from an
+            RLHF checkpoint, set `--resume_from_checkpoint` to the RLHF checkpoint and `--ref_adapters` to the SFT
+            checkpoint. Defaults to an empty list.
+        ref_model_type (Optional[str]): The model type of the reference model. Same as `model_type`. Defaults to None.
+        ref_model_revision (Optional[str]): The model revision of the reference model. Same as `model_revision`.
+            Defaults to None.
+        beta (Optional[float]): The beta parameter for RLHF, controlling the deviation from the reference model.
+            A higher value implies less deviation. If None, uses algorithm-specific defaults: 2.0 for 'simpo', 0.04
+            for 'grpo', 0.5 for 'gkd', and 0.1 for others. Defaults to None.
+        label_smoothing (float): The label smoothing value for DPO. A value of 0 disables it. Defaults to 0.
+        max_completion_length (int): The maximum generation length for GRPO/PPO/GKD algorithms. Defaults to 512.
+        loss_scale (Optional[str]): Overrides the template parameter. During RLHF training, this defaults to
+            'last_round'.
+        rpo_alpha (Optional[float]): The alpha parameter from the RPO paper, controlling the weight of the SFT loss
+            (NLL term). The loss is calculated as `dpo_loss + rpo_alpha * sft_loss`. If None, the SFT loss is not
+            included.
+        ld_alpha (Optional[float]): The alpha parameter from the LD-DPO paper, which weights the log probabilities of
+            the sequence part beyond the common prefix to mitigate length preference. Defaults to None.
+        discopop_tau (float): The temperature parameter from the DiscoPOP paper, used to scale the log-ratio. Effective
+            when `loss_type` is 'discopop'. Defaults to 0.05.
+        loss_type (Optional[List[str]]): The type of loss function. Defaults to algorithm-specific values (e.g.,
+            'sigmoid' for DPO). Multiple values can be passed for mixed training (MPO), which requires `loss_weights`
+            to be set.
+        loss_weights (Optional[List[float]]): When multiple `loss_type` values are set for DPO, this specifies the
+            weights for each loss term. Defaults to None.
+        cpo_alpha (float): The coefficient for the NLL loss in the CPO/SimPO loss function. Defaults to 1.0.
+        simpo_gamma (float): The reward margin term in the SimPO algorithm. The paper suggests a value between 0.5 and
+            1.5. Defaults to 1.0.
+        desirable_weight (float): In KTO, the weight applied to the desirable loss to counteract data imbalance.
+            Defaults to 1.0.
+        undesirable_weight (float): In KTO, the weight applied to the undesirable loss to counteract data imbalance.
+            Defaults to 1.0.
+        temperature (float): The temperature for sampling, used in PPO, GRPO, and GKD algorithms. Defaults to 0.9.
+        center_rewards_coefficient (Optional[float]): Used for Reward Model (RM) training. A coefficient to encourage
+            the reward model to output rewards with a mean of zero. A value of 0.01 is recommended. Defaults to None.
+        sft_alpha (float): The weight for the SFT loss component in GKD. The final loss is calculated as
+            gkd_loss + sft_alpha * sft_loss`. Defaults to 0.
+        lmbda (float): The lambda parameter for GKD, balancing policy and value losses. Defaults to 0.5.
+        seq_kd (bool): Deprecated. Sequential KD (teacher-generated responses) is not implemented.
+        gkd_logits_topk (Optional[int]): The number of top-k logits to use for KL divergence computation in GKD.
+            If None, uses full vocabulary for KL computation (more accurate but memory-intensive).
+            If set to a positive integer, only top-k teacher logits are used (more efficient).
+            When using `teacher_model_server`, this is limited by the server's `max_logprobs` setting
+            (vLLM default is 20, can be increased with `--max-logprobs`). Defaults to None.
+        max_new_tokens (Optional[int]): A backward-compatibility argument. Please use `max_completion_length` instead.
+            Defaults to None.
+    """
+    rlhf_type: Literal['dpo', 'orpo', 'simpo', 'kto', 'cpo', 'rm', 'ppo', 'grpo', 'gkd'] = 'dpo'
+    ref_model: Optional[str] = None
+    ref_adapters: List[str] = field(default_factory=list)
+    ref_model_type: Optional[str] = field(
+        default=None, metadata={'help': f'model_type choices: {list(MODEL_MAPPING.keys())}'})
+    ref_model_revision: Optional[str] = None
+
+    beta: Optional[float] = None
+    label_smoothing: float = 0
+    max_completion_length: int = 512
+    loss_scale: Optional[str] = None  # 'last_round'
+    # DPO
+    rpo_alpha: Optional[float] = None
+    ld_alpha: Optional[float] = None  # α parameter from the LD-DPO paper
+    discopop_tau: float = 0.05  # τ/temperature parameter from the DiscoPOP paper
+    loss_type: Optional[List[str]] = None
+    loss_weights: Optional[List[float]] = None
+    # CPO
+    cpo_alpha: float = 1.
+    # SimPO
+    simpo_gamma: float = 1
+    # KTO
+    desirable_weight: float = 1.0
+    undesirable_weight: float = 1.0
+    # PPO/GRPO/GKD
+    temperature: float = 0.9
+    # RM
+    center_rewards_coefficient: Optional[float] = None
+    # GKD
+    sft_alpha: float = 0
+    lmbda: float = 0.5
+    seq_kd: bool = False  # Deprecated
+    gkd_logits_topk: Optional[int] = None
+    # compat
+    max_new_tokens: Optional[int] = None  # use max_completion_length instead
+
+    def _prepare_training_args(self, training_args: Dict[str, Any]) -> None:
+        if self.rlhf_type == 'ppo':
+            training_args['world_size'] = self.global_world_size
+
+    def __post_init__(self):
+        self._process_loss_type()
+        self._init_grpo()
+        self._init_rm()
+        self._init_simpo()
+        self._init_max_completion_length()
+        self._init_padding_side()
+        self._set_default()
+        self._init_rollout()
+        self._init_teacher_deepspeed()
+        GRPOArguments.__post_init__(self)
+        SftArguments.__post_init__(self)
+        self._check_sequence_parallel()
+        self._check_teacher()
+        self._check_grpo()
+        self._check_gkd()
+
+        if isinstance(self.ref_adapters, str):
+            self.ref_adapters = [self.ref_adapters]
+        if self.rlhf_type == 'grpo' and self.beta == 0.0:
+            self.ref_model = None
+        elif self.rlhf_type in ['dpo', 'kto', 'ppo', 'grpo'] and self.tuner_type == 'full':
+            self.ref_model = self.ref_model or self.model
+            self.ref_model_type = self.ref_model_type or self.model_type
+            self.ref_model_revision = self.ref_model_revision or self.model_revision
+        elif self.ref_model is not None:
+            raise ValueError('CPO/ORPO or LoRA training does not require a ref_model to be passed in.')
+
+    def _set_loss_scale(self):
+        if self.loss_scale is None:
+            if self.rlhf_type == 'orpo' and not self.model_meta.is_multimodal:
+                # Avoid padding labels during the model's forward pass in multimodal models.
+                # Some multimodal models do not expand the image pad token.
+                self.loss_scale = 'default'
+            elif self.rlhf_type in ('grpo', 'gkd'):
+                if self.multi_turn_scheduler:
+                    self.loss_scale = 'default'
+                else:
+                    self.loss_scale = 'last_round'
+            else:
+                self.loss_scale = 'last_round'
+        super()._set_loss_scale()
+
+    def _process_loss_type(self):
+        if self.loss_type is None:
+            return
+
+        if isinstance(self.loss_type, list):
+            num_loss_types = len(self.loss_type)
+            if num_loss_types > 1:
+                assert self.rlhf_type == 'dpo', (f'Multiple loss types ({self.loss_type}) are only supported for DPO. '
+                                                 f'Current rlhf_type: {self.rlhf_type}.')
+                from trl.trainer.dpo_config import DPOConfig
+                assert 'loss_weights' in DPOConfig.__dict__, (
+                    'Multiple loss types requires trl >= 0.20, please install trl `pip install -U trl`')
+
+        if hasattr(self.loss_type, '__len__') and len(self.loss_type) == 1:
+            self.loss_type = self.loss_type[0]
+
+        # Validate loss_type
+        if self.loss_weights is not None:
+            assert self.rlhf_type == 'dpo'
+            loss_types = self.loss_type if isinstance(self.loss_type, list) else [self.loss_type]
+            if len(self.loss_weights) != len(loss_types):
+                raise ValueError(f'Length of loss_weights list ({self.loss_weights}) must match number of loss types '
+                                 f'({loss_types}).')
+
+    def _init_grpo(self):
+        if self.rlhf_type != 'grpo':
+            return
+        if self.cached_dataset or self.cached_val_dataset:
+            raise ValueError('cached_dataset is not supported for GRPO.')
+        if self.use_vllm:
+            set_default_ddp_config()
+        if self.async_generate or not self.use_vllm or self.vllm_mode == 'server':
+            self.sleep_level = 0
+        self.remove_unused_columns = False
+        logger.info(f'Setting args.remove_unused_columns: {self.remove_unused_columns}')
+        if self.truncation_strategy is None:
+            self.truncation_strategy = 'left'
+        if self.truncation_strategy not in {'left', 'delete'}:
+            raise ValueError("GRPO requires `truncation_strategy 'left' or 'delete'`, "
+                             f"Current value: `truncation_strategy='{self.truncation_strategy}'`.")
+        if self.beta is None:
+            self.beta = 0.04  # https://arxiv.org/abs/2402.03300
+        if self.async_generate:
+            logger.info('Using async mode. This is a approximate version which '
+                        'will use the old weights to generate responses to accelerate. '
+                        'This will ignore the `CLIP` of advantages, if you found the training '
+                        'is unstable, you may consider using --async_generate false.')
+        if 'soft_overlong' in self.reward_funcs:
+            assert self.soft_cache_length is not None, \
+                'The soft_cache_length must be set when using soft overlong rewards.'
+            if self.soft_max_length is None:
+                self.soft_max_length = self.max_completion_length
+                logger.info(f'Auto-configured soft_max_length = max_completion_length {self.max_completion_length}')
+
+        if self.kl_in_reward is None:
+            if self.advantage_estimator == 'grpo':
+                self.kl_in_reward = False
+            elif self.advantage_estimator in ['rloo', 'reinforce_plus_plus']:
+                self.kl_in_reward = True
+            else:
+                raise ValueError(f'Invalid advantage_estimator: {self.advantage_estimator}')
+
+        # disable normalization, REAL https://arxiv.org/abs/2602.05630
+        if self.loss_type == 'real':
+            self.scale_rewards = 'none'
+            logger.warning(
+                f"[REAL] scale_rewards='{self.scale_rewards}' is ignored. "
+                "It will be forced to 'none' because 'loss_type = real' does not support reward normalization.")
+
+        if self.scale_rewards is None:
+            if self.advantage_estimator == 'grpo':
+                self.scale_rewards = 'group'
+            elif self.advantage_estimator == 'rloo':
+                self.scale_rewards = 'none'
+            elif self.advantage_estimator == 'reinforce_plus_plus':
+                self.scale_rewards = 'batch'
+            else:
+                raise ValueError(f'Invalid advantage_estimator: {self.advantage_estimator}')
+
+    def _check_teacher(self):
+        self._teacher_use_disable_adapter = False
+
+        if self.rlhf_type not in ['grpo', 'gkd']:
+            if self.teacher_model is not None or self.teacher_model_server is not None:
+                logger.warning(f'teacher_model / teacher_model_server is ignored for rlhf_type={self.rlhf_type!r} '
+                               '(only used by gkd and grpo/OPD-RL).')
+            return
+        teacher_set = self.teacher_model is not None or self.teacher_model_server is not None
+        if not teacher_set:
+            if self.rlhf_type == 'gkd':
+                logger.info('No teacher_model specified. Using self-distillation mode (teacher = student).')
+                if self.use_liger_kernel:
+                    raise ValueError('Self-distillation mode with liger kernel loss is not supported yet')
+            if self.rlhf_type == 'grpo' and self.num_generations == 1:
+                raise ValueError('num_generations must be greater than 1 for GRPO')
+            return
+
+        if self.rlhf_type == 'grpo' and self.use_liger_kernel:
+            raise ValueError('OPD-RL is not supported with use_liger_kernel.')
+
+        if self.teacher_model is not None and self.teacher_model_server is not None:
+            raise ValueError('setting both `teacher_model` and `teacher_model_server` is not supported.')
+
+        # Validate teacher_model_server: accept single URL or JSON multi-teacher config.
+        if self.teacher_model_server is not None:
+            from swift.rlhf_trainers.gkd_helpers import parse_teacher_model_server
+
+            # Parse early to fail fast on invalid JSON; result is re-parsed by the trainer.
+            parse_teacher_model_server(self.teacher_model_server)
+
+        # Self-distillation: teacher_model == student model
+        if self.teacher_model is not None and self.teacher_model == self.model:
+            if self.tuner_type == 'lora':
+                logger.info('LoRA + same teacher_model: using disable_adapter() for fixed teacher (no extra model).')
+                self._teacher_use_disable_adapter = True
+                self.teacher_model = None
+            else:
+                # Full training + same teacher_model: a separate frozen copy will be loaded as fixed teacher.
+                pass
+
+    def _init_rollout(self):
+        if self.rlhf_type not in rlhf_support_vllm_types:
+            return
+
+        if self.use_vllm and os.getenv('SWIFT_AUDIO_LOAD_BACKEND') is None:
+            # align with vLLM audio load backend
+            os.environ['SWIFT_AUDIO_LOAD_BACKEND'] = 'soundfile_pyav'
+
+        if self.vllm_mode is not None and not self.use_vllm:
+            raise ValueError('vllm_mode is not supported when use_vllm is false')
+
+        if self.vllm_mode is None and self.use_vllm:
+            raise ValueError('vllm_mode is required when use_vllm is true')
+
+        self._init_external_vllm()
+
+        if self.vllm_mode == 'server':
+            assert not self.use_vllm or self.vllm_server_host is not None or self.vllm_server_base_url is not None
+
+        if self.async_generate:
+            assert self.vllm_mode == 'server', 'async generate require vllm_mode == server, '
+            'please deploy vLLM server by `swift rollout` and assign with `vllm_server_host` '
+            'for more infomations, please check '
+            'https://swift.readthedocs.io/en/latest/Instruction/GRPO/getstarted/GRPO.html'
+
+        if not self.use_vllm and self.vllm_tensor_parallel_size != 1:
+            self.vllm_tensor_parallel_size = 1
+            logger.warning('set vllm_tensor_parallel_size to 1 since use_vllm false')
+        self._external_vllm_warning()
+
+    def _init_padding_side(self):
+        if self.rlhf_type in {'ppo', 'gkd'}:
+            self.padding_side = 'left'
+            # TODO: streaming, MLLM
+
+    def _init_max_completion_length(self):
+        max_completion_length = self.response_length or self.max_new_tokens or self.max_completion_length
+        self.max_completion_length = self.max_new_tokens = self.response_length = max_completion_length
+
+    def _init_metric_for_best_model(self):
+        if self.rlhf_type == 'grpo' and self.metric_for_best_model is None:
+            self.metric_for_best_model = 'reward'
+        super()._init_metric_for_best_model()
+        if self.rlhf_type == 'ppo':
+            self.metric_for_best_model = None
+            self.greater_is_better = None
+
+    def _init_simpo(self):
+        if self.rlhf_type != 'simpo':
+            return
+
+        self.rlhf_type = 'cpo'
+        if self.loss_type is None:
+            self.loss_type = 'simpo'
+        if self.beta is None:
+            self.beta = 2.
+
+    def _init_rm(self):
+        if self.rlhf_type == 'rm':
+            self.task_type = 'seq_cls'
+            self.num_labels = 1
+
+    def _init_external_vllm(self):
+        if self.rlhf_type not in rlhf_support_vllm_types or (self.vllm_server_host is None
+                                                             and self.vllm_server_base_url is None):
+            return
+        from swift.rlhf_trainers import VLLMClient
+        if is_master():
+            logger.info('Start connecting to vLLM server')
+            self.vllm_client = VLLMClient(
+                base_urls=self.vllm_server_base_url,
+                hosts=self.vllm_server_host,
+                server_ports=self.vllm_server_port,
+                group_ports=self.vllm_server_group_port,
+                connection_timeout=self.vllm_server_timeout)
+            self.vllm_client.close_communicator()
+            self.vllm_client.init_communicator(device=get_current_device())
+            logger.info('Connected to vLLM server')
+
+    def _set_default(self):
+        if self.beta is None:
+            if self.rlhf_type == 'gkd':
+                self.beta = 0.5
+            else:
+                self.beta = 0.1
+        if self.loss_type is None:
+            if self.rlhf_type in ['dpo', 'cpo']:
+                self.loss_type = 'sigmoid'  # else None
+            elif self.rlhf_type in ['kto']:
+                self.loss_type = 'kto'
+            elif self.rlhf_type == 'grpo':
+                self.loss_type = 'grpo'
+        if self.gradient_accumulation_steps is None:
+            if self.rlhf_type == 'grpo':
+                self.gradient_accumulation_steps = 1
+                logger.info('Setting default gradient_accumulation_steps to 1 for GRPO.')
+
+    def _check_grpo(self):
+        if self.rlhf_type != 'grpo':
+            return
+        import importlib.metadata
+        import trl
+        from packaging import version
+        trl_version = version.parse(trl.__version__)
+        assert trl_version >= version.parse('0.20'), ('Your current version of `trl` is outdated. '
+                                                      'Please update it by running: pip install -U trl')
+        if is_mp() and self.use_vllm:
+            raise ValueError('GRPO with vLLM is not compatible with `device_map`. '
+                             'Please set NPROC_PER_NODE equal to num_processes.')
+        if self.use_liger_kernel:
+            liger_kernel_version = version.parse(importlib.metadata.version('liger-kernel'))
+            if liger_kernel_version < version.parse('0.7.0'):
+                raise ValueError('Please update liger-kernel to 0.7.0 or later: pip install -U liger-kernel')
+            if self.delta is not None:
+                raise ValueError('Liger loss does not support two-sided GRPO loss yet.')
+            if self.sequence_parallel_size > 1:
+                raise ValueError('Liger loss does not support sequence parallel yet.')
+            if self.padding_free:
+                raise ValueError('Liger loss does not support padding free yet.')
+            if self.top_entropy_quantile < 1.0:
+                raise ValueError('Liger loss does not support entropy mask yet.')
+            if self.log_entropy:
+                raise ValueError('Liger loss does not support log entropy yet.')
+            if self.off_policy_sequence_mask_delta is not None:
+                raise ValueError('Liger loss does not support off-policy sequence masking yet.')
+            assert self.importance_sampling_level in [
+                'token', 'sequence'
+            ], ('Liger loss currently only support token-level and sequence-level importance sampling. '
+                'Please set `importance_sampling_level` to `token` or `sequence`.')
+            if self.advantage_estimator != 'grpo':
+                raise ValueError('Liger loss currently only support grpo advantage estimator')
+
+        if self.async_generate and self.multi_turn_scheduler is not None:
+            raise NotImplementedError('Currently, async_generate is not supported with multi-turn functionality.')
+
+        self._check_opd_rl()
+
+    def _check_opd_rl(self):
+        """Fail-fast OPD-RL (teacher distillation on GRPO) parameter compatibility.
+
+        A teacher turns GRPO into OPD-RL, where the teacher signal is a *per-token* advantage
+        (the signed teacher log-ratio). Features that require a *per-sequence* advantage (typically
+        sign-based judgments) or reward variance are incompatible; reject them here rather than
+        deep inside the loss / advantage code. ``_check_teacher`` has already run, so
+        ``_teacher_use_disable_adapter`` is resolved.
+        """
+        opd_rl = (
+            self.teacher_model is not None or self.teacher_model_server is not None
+            or self._teacher_use_disable_adapter)
+        if not opd_rl:
+            return
+        # loss types / masks that reduce the advantage to a per-sequence scalar (sign-based).
+        if self.loss_type in ['real', 'fipo']:
+            raise ValueError(f'OPD-RL (teacher) does not support loss_type={self.loss_type!r} '
+                             '(it needs a per-sequence advantage).')
+        if self.off_policy_sequence_mask_delta is not None:
+            raise ValueError('OPD-RL (teacher) does not support off_policy_sequence_mask_delta '
+                             '(it needs a per-sequence advantage).')
+        # Pure distillation (no reward functions): the base GRPO advantage is 0, so reward-variance
+        # driven features have no signal to act on.
+        if not self.reward_funcs:
+            if self.dynamic_sample:
+                raise ValueError('dynamic_sample requires reward_funcs (it filters groups by reward std); '
+                                 'pure OPD-RL distillation has no reward variance.')
+            if self.scale_rewards == 'gdpo':
+                raise ValueError("scale_rewards='gdpo' requires reward_funcs; pure OPD-RL distillation has none.")
+
+    def _external_vllm_warning(self):
+        if self.rlhf_type not in rlhf_support_vllm_types or not self.vllm_server_host:
+            return
+
+        if self.vllm_max_model_len is not None:
+            logger.warning(
+                "Configuration conflict: 'vllm_max_model_len=%s' is ignored for external vLLM. "
+                'Please specify it when launching the inference service: '
+                '`swift rollout --vllm_max_model_len <value>`', self.vllm_max_model_len)
+
+    def _check_padding_free(self):
+        super()._check_padding_free()
+        if self.padding_free or self.packing:
+            supported_types = ['grpo', 'dpo', 'kto', 'gkd']
+            if self.rlhf_type not in supported_types:
+                raise NotImplementedError(
+                    f"The current rlhf_type '{self.rlhf_type}' does not support padding_free/packing. "
+                    'Please set --padding_free/packing to false.')
+
+    def _check_sequence_parallel(self):
+        if self.sequence_parallel_size > 1:
+            supported_types = ['grpo', 'dpo']
+            if self.rlhf_type not in supported_types:
+                raise NotImplementedError(
+                    f"The current rlhf_type '{self.rlhf_type}' does not support sequence_parallel. "
+                    'Please set --sequence_parallel_size to 1.')
+
+    def _init_teacher_deepspeed(self):
+        """Initialize teacher_deepspeed configuration similar to _init_deepspeed in SftArguments"""
+        if not self.teacher_deepspeed:
+            return
+
+        # Get the same ds_config_folder as main model
+        ds_config_folder = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'config'))
+        deepspeed_mapping = {
+            name: f'{name}.json'
+            for name in ['zero0', 'zero1', 'zero2', 'zero3', 'zero2_offload', 'zero3_offload']
+        }
+
+        # Check if teacher_deepspeed is a predefined name
+        for ds_name, ds_config in deepspeed_mapping.items():
+            if self.teacher_deepspeed == ds_name:
+                self.teacher_deepspeed = os.path.join(ds_config_folder, ds_config)
+                break
+
+        # Parse the config file to dict
+        self.teacher_deepspeed = json_parse_to_dict(self.teacher_deepspeed)
+        logger.info(f'Using teacher_deepspeed config: {self.teacher_deepspeed}')
+
+    def _check_gkd(self):
+        if self.rlhf_type != 'gkd':
+            return
+        if is_mp() and self.use_vllm:
+            raise ValueError('GKD with vLLM is not compatible with `device_map`. '
+                             'Please set NPROC_PER_NODE equal to num_processes.')
+
+        if self.async_generate:
+            raise NotImplementedError('Currently, async_generate is not supported for GKD.')
+
+        # seq_kd (teacher-generated responses) is not implemented; raise early.
+        if self.seq_kd:
+            raise NotImplementedError('seq_kd=True (Sequential KD with teacher generation) is deprecated.')
+
+        # When using teacher_model_server, gkd_logits_topk is required (API only returns top-k logprobs)
+        if self.teacher_model_server is not None:
+            if self.gkd_logits_topk is None:
+                raise ValueError('gkd_logits_topk is required when using teacher_model_server')
+
+        # Validate gkd_logits_topk
+        if self.gkd_logits_topk is not None and self.gkd_logits_topk <= 0:
+            raise ValueError(f'gkd_logits_topk must be a positive integer, got {self.gkd_logits_topk}')
+
+        if self.gkd_logits_topk is not None and self.use_liger_kernel:
+            raise ValueError('gkd_logits_topk is not supported when using liger kernel')

swift/arguments/sampling_args.py

@@ -0,0 +1,121 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import dataclasses
+import json
+from dataclasses import dataclass
+from datetime import datetime
+from typing import List, Literal, Optional
+
+from swift.utils import get_logger
+from .base_args import BaseArguments
+
+logger = get_logger()
+
+
+@dataclass
+class SamplingArguments(BaseArguments):
+    """A dataclass for configuring sampling parameters.
+
+    Args:
+        prm_model (Optional[str]): The type of the Process Reward Model (PRM). Can be a model ID (loaded via
+            'transformers' engine) or a PRM key defined in a plugin for custom inference. Defaults to None.
+        orm_model (Optional[str]): The type of the Outcome Reward Model (ORM). Typically a wildcard or test case,
+            usually defined in a plugin. Defaults to None.
+        sampler_type (Literal['sample', 'distill']): The type of sampling to perform. Supported types are 'sample' and
+            'distill'. Defaults to 'sample'.
+        sampler_engine (Literal['transformers', 'lmdeploy', 'vllm', 'no', 'client']): The inference engine for the
+            sampling model. Supported options are 'transformers', 'lmdeploy', 'vllm', 'client', and 'no'.
+            Defaults to 'transformers'.
+        output_dir (str): The directory to save the output files. Defaults to 'sample_output'.
+        output_file (Optional[str]): The name of the output file. If None, a timestamp will be used as the filename.
+            The path should not be included, only the filename. Only the '.jsonl' format is supported. Defaults to
+            None.
+        resume (bool): Whether to resume file. Defaults to False.
+        override_exist_file (bool): Whether to override the output file if it already exists. This is only effective
+            when `output_file` is specified. Defaults to False.
+        num_return_sequences (int): The number of raw sequences to return from sampling. Effective for the 'sample'
+            `sampler_type`. Defaults to 64.
+        num_sampling_batch_size (int): The batch size for each sampling iteration. Defaults to 1.
+        num_sampling_batches (Optional[int]): The total number of batches to sample. Defaults to None.
+        n_best_to_keep (int): The number of best sequences to keep after evaluation. Defaults to 5.
+        data_range (List[int]): Specifies the data shard to process. A list of two integers `[shard_index,
+            num_shards]`. For example, `[1, 3]` means the dataset is split into 3 shards and this process handles the
+            second shard (0-indexed). Defaults to [].
+        temperature (float): The temperature for sampling. Defaults to 1.0.
+        prm_threshold (float): The threshold for the Process Reward Model (PRM). Results with a score below this
+            threshold will be filtered out. Defaults to 0.0.
+        easy_query_threshold (Optional[float]): For a single query, if the proportion of correctly sampled sequences
+            (as evaluated by the ORM) is greater than this threshold, the query will be discarded. This prevents overly
+            simple queries from appearing in the final results. Defaults to None, which disables this filter.
+        engine_kwargs (Optional[str]): Additional arguments to pass to the `sampler_engine`, provided as a JSON string.
+            For example: '{"cache_max_entry_count":0.7}'. Defaults to None.
+        cache_files (List[str]): A list of cache files for a two-step sampling process to avoid OOM errors.
+            Step 1: Set `prm_model`, and `orm_model` to None. All generated sequences are saved to a file.
+            Step 2: Set `sampler_engine` to 'no' and provide the output file from Step 1 to `cache_files`.
+            This run will perform PRM and ORM evaluation on the cached results.
+            Note: The `--dataset` argument must still be provided, as IDs in the cache files are MD5 hashes of the
+            original data and need to be linked.
+    """
+    # rm models
+    prm_model: Optional[str] = None
+    orm_model: Optional[str] = None
+
+    # sampler settings
+    sampler_type: Literal['sample', 'distill'] = 'sample'
+    sampler_engine: Literal['transformers', 'lmdeploy', 'vllm', 'no', 'client'] = 'transformers'
+    output_dir: str = 'sample_output'
+    output_file: Optional[str] = None
+    resume: bool = False
+    override_exist_file: bool = False
+    num_return_sequences: int = 64
+    num_sampling_batch_size: int = 1
+    num_sampling_batches: Optional[int] = None
+    n_best_to_keep: int = 5
+    data_range: List[int] = dataclasses.field(default_factory=list)
+
+    # generate settings
+    temperature: float = 1.0
+    prm_threshold: float = 0.0
+    easy_query_threshold: Optional[float] = None
+
+    # engine settings
+    engine_kwargs: Optional[str] = None
+
+    # Vanilla
+    cache_files: List[str] = dataclasses.field(default_factory=list)
+
+    def _init_model_info(self):
+        if self.sampler_engine != 'client':
+            return super()._init_model_info()
+        else:
+            self.model_info = None
+            self.model_meta = None
+        self.task_type = 'causal_lm'
+        return
+
+    def __post_init__(self):
+        if self.sampler_engine == 'pt':
+            self.sampler_engine = 'transformers'  # compat swift3.x
+        if self.output_file is None:
+            now = datetime.now()
+            formatted_time = now.strftime('%Y-%m-%d-%H-%M-%S')
+            self.output_file = formatted_time + '.jsonl'
+            logger.info(f'Setting output_file to {self.output_file}')
+        else:
+            if '/' in self.output_file or '\\' in self.output_file:
+                raise ValueError(f'Please use a string prefix without directory to '
+                                 f'`--output_file` but now is: {self.output_file}')
+        self.padding_side = 'left'
+        if self.engine_kwargs is not None:
+            self.engine_kwargs = json.loads(self.engine_kwargs)
+        else:
+            self.engine_kwargs = {}
+
+        super().__post_init__()
+
+        if self.system is not None:
+            self.system_message = [{
+                'role': 'system',
+                'content': self.system,
+            }]
+        else:
+            self.system_message = []

swift/arguments/sft_args.py

@@ -0,0 +1,425 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import os
+from dataclasses import dataclass
+from transformers.utils.versions import require_version
+from typing import Literal, Optional
+
+from swift.trainers import Seq2SeqTrainingArguments, TrainerFactory
+from swift.trainers.utils import prepare_deepspeed_elastic_config
+from swift.utils import (add_version_to_work_dir, get_device_count, get_logger, get_pai_tensorboard_dir, is_mp,
+                         is_pai_training_job, is_swanlab_available, json_parse_to_dict, to_abspath)
+from .base_args import BaseArguments
+from .tuner_args import TunerArguments
+
+logger = get_logger()
+
+
+@dataclass
+class SwanlabArguments:
+    """Arguments for configuring Swanlab for experiment result logging.
+
+    This dataclass stores all the configuration parameters required for initializing and using Swanlab to track
+    experiments.
+
+    Args:
+        swanlab_token (Optional[str]): The API key for SwanLab. You can also specify it using the `SWANLAB_API_KEY`
+            environment variable.
+        swanlab_project (str): The SwanLab project, which can be created in advance on the page
+            [https://swanlab.cn/space/~](https://swanlab.cn/space/~) or created automatically.
+            The default is "ms-swift".
+        swanlab_workspace (Optional[str]): The SwanLab workspace. Defaults to `None`, in which case the username
+            associated with the API key will be used.
+        swanlab_exp_name (Optional[str]): The name of the experiment. If `None`, it will default to the value of the
+            `output_dir` argument.
+        swanlab_notification_method (Optional[str]): The notification method for SwanLab when training completes
+            or errors occur. For details, refer to [here](https://docs.swanlab.cn/plugin/notification-dingtalk.html).
+            Supports 'dingtalk', 'lark', 'email', 'discord', 'wxwork', 'slack'.
+        swanlab_webhook_url (Optional[str]): Defaults to None. The webhook URL corresponding to
+            SwanLab's `swanlab_notification_method`.
+        swanlab_secret (Optional[str]): Defaults to None. The secret corresponding to
+            SwanLab's `swanlab_notification_method`.
+        swanlab_sender_email (Optional[str]): The email address of the sender. Required when
+            `swanlab_notification_method` is 'email'.
+        swanlab_receiver_email (Optional[str]): The email address of the receiver. Required when
+            `swanlab_notification_method` is 'email'.
+        swanlab_smtp_server (Optional[str]): The SMTP server address for email notification (e.g., 'smtp.qq.com').
+        swanlab_smtp_port (Optional[int]): The SMTP server port for email notification (e.g., 465).
+        swanlab_email_language (Optional[str]): email messages language. Supports 'zh', 'en'. The default is "zh".
+        swanlab_mode (Literal['cloud', 'local']): The operation mode, either 'cloud' for cloud-based logging or 'local'
+            for local-only logging.
+    """
+    swanlab_token: Optional[str] = None
+    swanlab_project: str = 'ms-swift'
+    swanlab_workspace: Optional[str] = None
+    swanlab_exp_name: Optional[str] = None
+    swanlab_notification_method: Optional[str] = None
+    swanlab_webhook_url: Optional[str] = None
+    swanlab_secret: Optional[str] = None
+    swanlab_sender_email: Optional[str] = None
+    swanlab_receiver_email: Optional[str] = None
+    swanlab_smtp_server: Optional[str] = None
+    swanlab_smtp_port: Optional[int] = None
+    swanlab_email_language: Optional[str] = 'zh'
+    swanlab_mode: Literal['cloud', 'local'] = 'cloud'
+
+    def _init_swanlab(self):
+        if not is_swanlab_available():
+            raise ValueError('You are using swanlab as `report_to`, please install swanlab by '
+                             '`pip install swanlab`')
+        if not self.swanlab_exp_name:
+            self.swanlab_exp_name = self.output_dir
+        import swanlab
+        from swanlab.integration.transformers import SwanLabCallback
+        from transformers.integrations import INTEGRATION_TO_CALLBACK
+        if self.swanlab_token:
+            swanlab.login(self.swanlab_token)
+
+        if self.swanlab_notification_method is not None:
+            from swanlab.plugin.notification import (DingTalkCallback, DiscordCallback, EmailCallback, LarkCallback,
+                                                     SlackCallback, WXWorkCallback)
+            notification_mapping = {
+                'lark': LarkCallback,
+                'dingtalk': DingTalkCallback,
+                'email': EmailCallback,
+                'discord': DiscordCallback,
+                'wxwork': WXWorkCallback,
+                'slack': SlackCallback,
+            }
+            callback_cls = notification_mapping.get(self.swanlab_notification_method)
+            if callback_cls is None:
+                raise ValueError(
+                    f'Unsupported swanlab_notification_method: "{self.swanlab_notification_method}". Supported methods'
+                    f' are: {list(notification_mapping.keys())}')
+
+            if self.swanlab_notification_method == 'email':
+                if not (self.swanlab_sender_email and self.swanlab_receiver_email and self.swanlab_smtp_server
+                        and self.swanlab_smtp_port):
+                    raise ValueError("When 'swanlab_notification_method' is 'email', both 'swanlab_sender_email' "
+                                     "and 'swanlab_receiver_email' and 'swanlab_smtp_server' and 'swanlab_smtp_port' "
+                                     'must be provided.')
+                callback = EmailCallback(
+                    sender_email=self.swanlab_sender_email,
+                    receiver_email=self.swanlab_receiver_email,
+                    password=self.swanlab_secret,
+                    smtp_server=self.swanlab_smtp_server,
+                    port=self.swanlab_smtp_port,
+                    language=self.swanlab_email_language)
+            else:
+                callback = callback_cls(
+                    webhook_url=self.swanlab_webhook_url,
+                    secret=self.swanlab_secret,
+                )
+            swanlab.register_callbacks([callback])
+
+        INTEGRATION_TO_CALLBACK['swanlab'] = SwanLabCallback(
+            project=self.swanlab_project,
+            workspace=self.swanlab_workspace,
+            experiment_name=self.swanlab_exp_name,
+            config={'UPPERFRAME': '🐦‍⬛ms-swift'},
+            mode=self.swanlab_mode,
+        )
+
+
+@dataclass
+class SftArguments(SwanlabArguments, TunerArguments, BaseArguments, Seq2SeqTrainingArguments):
+    """Arguments pertaining to the training process.
+
+    SftArguments is a dataclass that inherits from multiple argument classes: SwanlabArguments, TunerArguments,
+    BaseArguments, Seq2SeqTrainingArguments.
+
+    Args:
+        add_version (bool): Whether to add a versioned subdirectory like '<version>-<timestamp>' to the `output_dir` to
+            prevent overwriting existing checkpoints. Defaults to True.
+        create_checkpoint_symlink (bool): Whether to create additional symbolic links for checkpoints, which can be
+            useful for automated training scripts. The symlinks for the best and last models will be created at
+            `f'{output_dir}/best'` and `f'{output_dir}/last'`, respectively. Defaults to False.
+        output_dir (Optional[str]): The directory to save model outputs. Defaults to 'output/<model_name>'.
+        learning_rate (Optional[float]): The learning rate. Defaults to 1e-5 for full-parameter training and 1e-4 for
+            tuners like LoRA.
+            Note: To set a minimum learning rate (min_lr), you can pass the arguments
+            --lr_scheduler_type cosine_with_min_lr --lr_scheduler_kwargs '{"min_lr": 1e-6}'.
+        eval_strategy (Optional[str]): The evaluation strategy. By default, it aligns with `save_strategy`. It will
+            default to 'no' if no validation dataset is provided (i.e., `val_dataset` and `eval_dataset` are not used,
+            and `split_dataset_ratio` is 0).
+        fp16 (Optional[bool]): Defaults to None.
+        bf16 (Optional[bool]): Defaults to None.
+        max_new_tokens (int): Overrides generation parameters. The maximum number of new tokens to generate when
+            `predict_with_generate` is True. Defaults to 64.
+        temperature (float): Overrides generation parameters. The temperature for sampling when `predict_with_generate`
+            is True. Defaults to 0.0.
+        load_args (bool): Whether to load `args.json` from a saved directory when `--resume_from_checkpoint`,
+            `--model`, or `--adapters` is specified. For details on which keys are loaded, refer to `base_args.py`.
+            Defaults to `True` for inference and exporting, and `False` for training. This argument typically does not
+            need to be modified.
+        zero_hpz_partition_size (Optional[int]): A feature of ZeRO++. Enables model sharding within a node and data
+            sharding between nodes. If you encounter `grad_norm` NaN issues, consider trying `--torch_dtype float16`.
+            Defaults to None.
+        deepspeed_autotp_size (Optional[int]): The tensor parallelism size for DeepSpeed AutoTP. To use this, the
+            `--deepspeed` argument must be set to 'zero0', 'zero1', or 'zero2'. Note: This feature only supports
+            full-parameter fine-tuning. Defaults to None.
+    """
+    add_version: bool = True
+    create_checkpoint_symlink: bool = False
+
+    # override
+    output_dir: Optional[str] = None
+    learning_rate: Optional[float] = None
+    eval_strategy: Optional[str] = None  # steps, epoch
+    fp16: Optional[bool] = None
+    bf16: Optional[bool] = None
+
+    # extra
+    max_new_tokens: int = 64
+    temperature: float = 0.
+    load_args: bool = False
+
+    # zero++
+    zero_hpz_partition_size: Optional[int] = None
+
+    # auto_tp
+    deepspeed_autotp_size: Optional[int] = None
+
+    # fsdp
+    fsdp: Optional[str] = None
+
+    def _check_padding_free(self):
+        if self.padding_free or self.packing:
+            if self.packing:
+                feature = 'packing'
+                self.padding_free = True
+            else:
+                feature = 'padding_free'
+            supported_impls = ['flash_attn', 'flash_attention_2', 'flash_attention_3', 'flash_attention_4']
+            if self.attn_impl not in supported_impls:
+                supported_impls_str = ', '.join([f'"{impl}"' for impl in supported_impls])
+                raise ValueError(f'The "{feature}" feature requires a flash attention implementation. '
+                                 f'Please use one of: {supported_impls_str}.')
+
+    def __post_init__(self) -> None:
+        if self.resume_from_checkpoint:
+            self.resume_from_checkpoint = to_abspath(self.resume_from_checkpoint, True)
+            # The non-resume_only_model will have its weights loaded in the trainer.
+            if self.resume_only_model:
+                if self.tuner_type == 'full':
+                    self.model = self.resume_from_checkpoint
+                else:
+                    self.adapters = [self.resume_from_checkpoint]
+        BaseArguments.__post_init__(self)
+        self._init_override()
+        TunerArguments.__post_init__(self)
+        self._check_padding_free()
+        if self.vit_gradient_checkpointing is None:
+            self.vit_gradient_checkpointing = not self.freeze_vit
+        if self.optimizer is None:
+            if self.lorap_lr_ratio:
+                self.optimizer = 'lorap'
+            elif self.use_galore:
+                self.optimizer = 'galore'
+
+        if len(self.dataset) == 0 and len(self.cached_dataset) == 0:
+            raise ValueError(f'self.dataset: {self.dataset}, self.cached_dataset: {self.cached_dataset}. '
+                             'Please input the training dataset.')
+
+        self._handle_pai_compat()
+
+        self._init_deepspeed()
+        self._init_fsdp()
+        self._init_device()
+
+        if getattr(self, 'accelerator_config', None) is None:
+            self.accelerator_config = {'dispatch_batches': False}
+        if not (self.eval_dataset or self._val_dataset_exists):
+            self.eval_strategy = 'no'
+        self.training_args = TrainerFactory.get_training_args(self)
+        self.training_args.remove_unused_columns = False
+        self._add_version()
+
+        if 'swanlab' in self.report_to:
+            self._init_swanlab()
+
+    def _init_override(self):
+        self._init_output_dir()
+        self._init_metric()
+
+        if self.learning_rate is None:
+            if self.tuner_type == 'full':
+                self.learning_rate = 1e-5
+            else:
+                self.learning_rate = 1e-4
+        self._init_eval_strategy()
+
+    def _init_deepspeed(self):
+        if self.deepspeed:
+            require_version('deepspeed')
+            if is_mp() and not self.use_ray:
+                raise ValueError('DeepSpeed is not compatible with `device_map`. '
+                                 f'n_gpu: {get_device_count()}, '
+                                 f'local_world_size: {self.local_world_size}.')
+
+            ds_config_folder = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'config'))
+            deepspeed_mapping = {
+                name: f'{name}.json'
+                for name in ['zero0', 'zero1', 'zero2', 'zero3', 'zero2_offload', 'zero3_offload']
+            }
+            for ds_name, ds_config in deepspeed_mapping.items():
+                if self.deepspeed == ds_name:
+                    self.deepspeed = os.path.join(ds_config_folder, ds_config)
+                    break
+
+            self.deepspeed = json_parse_to_dict(self.deepspeed)
+            if self.zero_hpz_partition_size is not None:
+                assert 'zero_optimization' in self.deepspeed
+                self.deepspeed['zero_optimization']['zero_hpz_partition_size'] = self.zero_hpz_partition_size
+                logger.warn('If `zero_hpz_partition_size`(ZeRO++) causes grad_norm NaN, please'
+                            ' try `--torch_dtype float16`')
+            if self.deepspeed_autotp_size is not None:
+                assert self.deepspeed is not None, (
+                    'To use `deepspeed_autotp_size`, you need to additionally set the `--deepspeed` argument.')
+                self.deepspeed.setdefault('tensor_parallel', {})['autotp_size'] = self.deepspeed_autotp_size
+                self.deepspeed.setdefault('zero_optimization', {})['gather_16bit_weights_on_model_save'] = True
+            if 'deepspeed_elastic' in set(getattr(self, 'callbacks', []) or []):
+                prepare_deepspeed_elastic_config(self)
+            logger.info(f'Using deepspeed: {self.deepspeed}')
+
+    def _init_fsdp(self):
+        if not self.fsdp:
+            self.fsdp = []
+            return
+
+        if is_mp() and not self.use_ray:
+            raise ValueError('FSDP2 is not compatible with `device_map`. '
+                             f'n_gpu: {get_device_count()}, '
+                             f'local_world_size: {self.local_world_size}.')
+        if self.deepspeed:
+            raise ValueError('FSDP2 is not compatible with DeepSpeed.')
+
+        fsdp_config_folder = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'config'))
+
+        # FSDP2 preset configurations
+        fsdp_mapping = {
+            'fsdp2': 'fsdp2.json',
+        }
+
+        fsdp_config_path = self.fsdp
+        for fsdp_name, fsdp_config in fsdp_mapping.items():
+            if self.fsdp == fsdp_name:
+                fsdp_config_path = os.path.join(fsdp_config_folder, fsdp_config)
+                break
+
+        fsdp_config_dict = json_parse_to_dict(fsdp_config_path)
+
+        # Extract fsdp string options (e.g., "full_shard auto_wrap offload")
+        fsdp_options = fsdp_config_dict.get('fsdp', 'full_shard auto_wrap')
+        self.fsdp = fsdp_options
+
+        # Extract fsdp_config dict
+        self.fsdp_config = fsdp_config_dict.get('fsdp_config', {})
+
+        # Set FSDP_VERSION environment variable for accelerate to recognize FSDP2
+        fsdp_version = self.fsdp_config.get('fsdp_version', 2)
+        os.environ['FSDP_VERSION'] = str(fsdp_version)
+
+        # Set environment variable to optimize NCCL memory usage
+        if 'TORCH_NCCL_AVOID_RECORD_STREAMS' not in os.environ:
+            os.environ['TORCH_NCCL_AVOID_RECORD_STREAMS'] = '1'
+
+        # Check FSDP2 compatibility with other training arguments
+        self._check_fsdp2_compatibility()
+
+        logger.info(f'Using FSDP2: fsdp={self.fsdp}, fsdp_config={self.fsdp_config}')
+
+    def _check_fsdp2_compatibility(self):
+        """Check for incompatible argument combinations with FSDP2.
+
+        FSDP2 has several known limitations:
+        1. save_only_model=True + SHARDED_STATE_DICT: Can't save only model weights with sharded state dict
+        2. gradient_checkpointing=True: Should use activation_checkpointing in fsdp_config instead
+        """
+        state_dict_type = self.fsdp_config.get('state_dict_type', 'SHARDED_STATE_DICT')
+
+        # Check 1: save_only_model + SHARDED_STATE_DICT
+        if getattr(self, 'save_only_model', False) and 'SHARDED' in state_dict_type.upper():
+            raise ValueError(
+                'FSDP2 with SHARDED_STATE_DICT is not compatible with save_only_model=True. '
+                'Either set save_only_model=False, or change state_dict_type to FULL_STATE_DICT in fsdp_config. '
+                'Note: FULL_STATE_DICT requires more memory and is slower.')
+
+        # Check 2: gradient_checkpointing should be disabled, use activation_checkpointing instead
+        if getattr(self, 'gradient_checkpointing', False):
+            activation_checkpointing = self.fsdp_config.get('activation_checkpointing', False)
+            if activation_checkpointing:
+                logger.warning('Both gradient_checkpointing and fsdp_config.activation_checkpointing are enabled. '
+                               'For FSDP2, it is recommended to use only activation_checkpointing in fsdp_config. '
+                               'Disabling gradient_checkpointing automatically.')
+                self.gradient_checkpointing = False
+            else:
+                logger.warning(
+                    'gradient_checkpointing is enabled with FSDP2. '
+                    'For better performance, consider using activation_checkpointing in fsdp_config instead. '
+                    'Add "activation_checkpointing": true to your fsdp_config.')
+
+    def _handle_pai_compat(self) -> None:
+        if not is_pai_training_job():
+            return
+
+        logger.info('Handle pai compat...')
+        pai_tensorboard_dir = get_pai_tensorboard_dir()
+        if self.logging_dir is None and pai_tensorboard_dir is not None:
+            self.logging_dir = pai_tensorboard_dir
+            logger.info(f'Setting args.logging_dir: {self.logging_dir}')
+        self.add_version = False
+        logger.info(f'Setting args.add_version: {self.add_version}')
+
+    def _add_version(self):
+        """Prepare the output_dir"""
+        if self.add_version:
+            self.output_dir = add_version_to_work_dir(self.output_dir)
+            logger.info(f'output_dir: {self.output_dir}')
+
+        if self.logging_dir is None:
+            self.logging_dir = f'{self.output_dir}/runs'
+
+        self.logging_dir = to_abspath(self.logging_dir)
+        os.makedirs(self.output_dir, exist_ok=True)
+
+        if self.run_name is None:
+            self.run_name = self.output_dir
+
+        self.training_args.output_dir = self.output_dir
+        self.training_args.run_name = self.run_name
+        self.training_args.logging_dir = self.logging_dir
+
+    def _init_output_dir(self):
+        if self.output_dir is None:
+            self.output_dir = f'output/{self.model_suffix}'
+        self.output_dir = to_abspath(self.output_dir)
+
+    def _init_eval_strategy(self):
+        if self.eval_strategy is None:
+            self.eval_strategy = self.save_strategy
+        if self.eval_strategy == 'no':
+            self.eval_steps = None
+            if self.split_dataset_ratio > 0:
+                self.split_dataset_ratio = 0.
+                logger.info(f'Setting args.split_dataset_ratio: {self.split_dataset_ratio}')
+        elif self.eval_strategy == 'steps' and self.eval_steps is None:
+            self.eval_steps = self.save_steps
+        self.evaluation_strategy = self.eval_strategy
+
+    def _init_metric(self):
+        if self.eval_metric is None:
+            if self.task_type == 'causal_lm' and self.predict_with_generate:
+                self.eval_metric = 'nlg'
+            elif self.task_type == 'embedding':
+                self.eval_metric = 'infonce' if self.loss_type == 'infonce' else 'paired'
+            elif self.task_type in {'reranker', 'generative_reranker'}:
+                self.eval_metric = 'reranker'
+        if self.eval_metric == 'nlg':
+            require_version('jieba', 'Setting `--eval_metric nlg` requires installing the jieba dependency.')
+        self._init_metric_for_best_model()
+
+    def _init_metric_for_best_model(self):
+        if self.metric_for_best_model is None:
+            self.metric_for_best_model = 'rouge-l' if self.predict_with_generate else 'loss'
+        if self.greater_is_better is None and self.metric_for_best_model is not None:
+            self.greater_is_better = 'loss' not in self.metric_for_best_model

swift/arguments/tuner_args.py

@@ -0,0 +1,220 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from dataclasses import dataclass, field
+from transformers.utils import strtobool
+from typing import List, Literal, Optional
+
+from swift.utils import get_logger
+
+logger = get_logger()
+
+
+@dataclass
+class TunerArguments:
+    """
+    TunerArguments is a dataclass that holds configuration for various tuners.
+
+    Args:
+        freeze_parameters (List[str]): A list of prefixes for parameters that should be frozen during training.
+            Defaults to an empty list `[]`.
+        freeze_parameters_regex (Optional[str]): A regular expression to match the names of parameters that should be
+            frozen. Defaults to `None`.
+        freeze_parameters_ratio (float): The ratio of parameters to freeze, starting from the bottom layers upwards
+            (from 0.0 to 1.0). Setting this to 1.0 freezes all model parameters, which can be useful when selectively
+            unfreezing specific parameters with `trainable_parameters`. Defaults to 0.0.
+        trainable_parameters (List[str]): A list of prefixes for parameters that should be made explicitly trainable.
+            Defaults to an empty list `[]`.
+        trainable_parameters_regex (Optional[str]): A regular expression to match the names of parameters that should
+            be made explicitly trainable. Defaults to `None`.
+            Note on parameter freezing priority: The `trainable_*` arguments have higher priority than the `freeze_*`
+            arguments. The freezing logic is applied as follows:
+            Firstly, all parameters are set to trainable.
+            Then, `freeze_parameters`, `freeze_parameters_regex`, and `freeze_parameters_ratio` are applied to freeze
+            parts of the model.
+            Finally, `trainable_parameters` and `trainable_parameters_regex` are used to unfreeze specific parameters,
+            ensuring they are trainable regardless of the freezing rules.
+
+        freeze_llm (bool): For multi-modal models only. If `True`, it affects the Large Language Model (LLM) part.
+            In full fine-tuning, this freezes the LLM weights. In LoRA training with `target_modules=['all-linear']`,
+            this prevents adding LoRA modules to the LLM. Defaults to `False`.
+        freeze_vit (bool): For multi-modal models only. If `True`, it affects the Vision/Audio Transformer (ViT) part.
+            In full fine-tuning, this freezes the ViT weights. In LoRA training with `target_modules=['all-linear']`,
+            this prevents adding LoRA modules to the ViT. Note: 'vit' can refer to `vision_tower` and `audio_tower`.
+            Defaults to `True`.
+        freeze_aligner (bool): For multi-modal models only. If `True`, it affects the aligner (projector) part.
+            In full fine-tuning, this freezes the aligner weights. In LoRA training with
+            `target_modules=['all-linear']`, this prevents adding LoRA modules to the aligner. Defaults to `True`.
+
+        target_modules (List[str]): List of target modules for tuning. Default is ['all-linear'].
+        target_regex (Optional[str]): Regular expression to match target modules. Default is None.
+        target_parameters (Optional[List[str]]): A list of parameter names to be replaced by LoRA modules. This is
+            similar to `target_modules` but targets parameters directly, which is useful for layers like MoE that use
+            `nn.Parameter` instead of `nn.Linear`. Requires `peft>=0.17.0`. Defaults to `None`.
+        modules_to_save (List[str]): List of modules to save. Default is an empty list.
+
+        lora_rank (int): Rank for LoRA. Default is 8.
+        lora_alpha (int): Alpha value for LoRA. Default is 32.
+        lora_dropout (float): Dropout rate for LoRA. Default is 0.05.
+        lora_bias (Literal['none', 'all']): The possible values are 'none' and 'all'. If set to 'all', all biases
+            will be trainable. Default is 'none'.
+        lora_dtype (Literal): Data type for LoRA. Default is 'AUTO'. Allowed values are 'fp16', 'bf16', 'fp32', 'AUTO'.
+        lorap_lr_ratio (float): Learning rate ratio for LoRA. Default is None.
+        use_rslora (bool): Flag to indicate if RSLora is used. Default is False.
+        use_dora (bool): Flag to indicate if Dora is used. Default is False.
+
+        lora_ga_batch_size (int): Batch size used for estimating gradients during initialization in LoRA-GA. Default
+            value is 2.
+        lora_ga_iters (int): Number of iterations for estimating gradients during initialization in LoRA-GA. Default
+            value is 2.
+        lora_ga_max_length (int): Maximum input length for estimating gradients during initialization in LoRA-GA.
+            Default value is 1024.
+        lora_ga_direction (str): Initial direction used for gradient estimation during initialization in LoRA-GA.
+            Default value is `ArB2r`. Allowed: `ArBr`, `A2rBr`, `ArB2r`, and `random`.
+        lora_ga_scale (str): The scaling method for initialization in LoRA-GA.
+            Default value is `stable`. Allowed values are: `gd`, `unit`, `stable`, and `weightS`.
+        lora_ga_stable_gamma (int): The gamma value when choosing `stable` scaling for initialization. Default
+            value is 16.
+
+        init_weights (str): The method for initializing adapter weights. For LoRA, options include 'true', 'false',
+            'gaussian', 'pissa', 'pissa_niter_[number of iters]', 'olora', 'loftq', and 'lora-ga'. For BoNE,
+            options are 'true', 'false', and 'bat'. Defaults to 'true'.
+
+        fourier_n_frequency (int): Number of frequencies for FourierFT. Default is 2000.
+        fourier_scaling (float): Scaling factor for FourierFT. Default is 300.0.
+
+        boft_block_size (int): Block size for BOFT. Default is 4.
+        boft_block_num (int): Number of blocks for BOFT. Default is 0.
+        boft_n_butterfly_factor (int): Butterfly factor for BOFT. Default is 1.
+        boft_dropout (float): Dropout rate for BOFT. Default is 0.0.
+
+        vera_rank (int): Rank for Vera. Default is 256.
+        vera_projection_prng_key (int): PRNG key for Vera projection. Default is 0.
+        vera_dropout (float): Dropout rate for Vera. Default is 0.0.
+        vera_d_initial (float): Initial value for Vera D. Default is 0.1.
+
+        adapter_act (str): Activation function for adapter. Default is 'gelu'.
+        adapter_length (int): Length of the adapter. Default is 128.
+
+        adalora_target_r (int): Target rank for AdaLoRA. Default is 8.
+        adalora_init_r (int): Initial rank for AdaLoRA. Default is 12.
+        adalora_tinit (int): Initial T value for AdaLoRA. Default is 100.
+        adalora_tfinal (int): Final T value for AdaLoRA. Default is 1000.
+        adalora_deltaT (int): Delta T value for AdaLoRA. Default is 10.
+        adalora_beta1 (float): Beta1 value for AdaLoRA. Default is 0.85.
+        adalora_beta2 (float): Beta2 value for AdaLoRA. Default is 0.85.
+        adalora_orth_reg_weight (float): Orthogonal regularization weight for AdaLoRA. Default is 0.5.
+
+        llamapro_num_new_blocks (int): Number of new blocks for LLaMAPro. Default is 4.
+        llamapro_num_groups (Optional[int]): Number of groups for LLaMAPro. Default is None.
+
+        reft_layer_key (Optional[str]): Key identifier for ReFT layer. Default is None.
+        reft_layers (Optional[List[int]]): List of layers involved in ReFT. Default is None.
+        reft_rank (int): Rank parameter for ReFT. Default is 4.
+        reft_intervention_type (Literal): Type of intervention for ReFT. Default is 'LoreftIntervention'.
+        reft_args (Optional[str]): Additional arguments for ReFT. Default is None.
+    """
+    # full
+    freeze_parameters: List[str] = field(default_factory=list)
+    freeze_parameters_regex: Optional[str] = None
+    freeze_parameters_ratio: float = 0.  # 0 ~ 1
+    trainable_parameters: List[str] = field(default_factory=list)
+    trainable_parameters_regex: Optional[str] = None
+    # lora or full
+    freeze_llm: bool = False
+    freeze_vit: bool = True
+    freeze_aligner: bool = True
+    # tuners
+    target_modules: List[str] = field(default_factory=lambda: ['all-linear'])
+    target_regex: Optional[str] = None
+    target_parameters: Optional[List[str]] = None
+    # e.g. ['wte', 'ln_1', 'ln_2', 'ln_f', 'lm_head']
+    modules_to_save: List[str] = field(default_factory=list)
+
+    # lora
+    lora_rank: int = 8
+    lora_alpha: int = 32
+    lora_dropout: float = 0.05
+    lora_bias: Literal['none', 'all'] = 'none'
+    lora_dtype: Literal['float16', 'bfloat16', 'float32', None] = None
+    lorap_lr_ratio: Optional[float] = None
+    use_rslora: bool = False
+    use_dora: bool = False
+
+    # lora_ga
+    lora_ga_batch_size: int = 2
+    lora_ga_iters: int = 2
+    lora_ga_max_length: int = 1024
+    lora_ga_direction: str = 'ArB2r'
+    lora_ga_scale: str = 'stable'
+    lora_ga_stable_gamma: int = 16
+
+    # Lora: Literal['gaussian', 'pissa', 'pissa_niter_[number of iters]', 'olora', 'loftq', 'true', 'false', 'lora-ga']
+    # Bone: Literal['bat', 'true', 'false']
+    init_weights: str = 'true'
+
+    # fourierft
+    fourier_n_frequency: int = 2000
+    fourier_scaling: float = 300.0
+
+    # BOFT
+    boft_block_size: int = 4
+    boft_block_num: int = 0
+    boft_n_butterfly_factor: int = 1
+    boft_dropout: float = 0.0
+
+    # Vera
+    vera_rank: int = 256
+    vera_projection_prng_key: int = 0
+    vera_dropout: float = 0.0
+    vera_d_initial: float = 0.1
+
+    # adapter
+    adapter_act: str = 'gelu'
+    adapter_length: int = 128
+
+    # adalora
+    adalora_target_r: int = 8
+    adalora_init_r: int = 12
+    adalora_tinit: int = 0
+    adalora_tfinal: int = 0
+    adalora_deltaT: int = 1
+    adalora_beta1: float = 0.85
+    adalora_beta2: float = 0.85
+    adalora_orth_reg_weight: float = 0.5
+
+    # llamapro
+    llamapro_num_new_blocks: int = 4
+    llamapro_num_groups: Optional[int] = None
+
+    # reft
+    reft_layer_key: Optional[str] = None
+    reft_layers: Optional[List[int]] = None
+    reft_rank: int = 4
+    reft_intervention_type: Literal['NoreftIntervention', 'LoreftIntervention', 'ConsreftIntervention',
+                                    'LobireftIntervention', 'DireftIntervention',
+                                    'NodireftIntervention'] = 'LoreftIntervention'
+    reft_args: Optional[str] = None
+
+    def __post_init__(self):
+        if isinstance(self.init_weights, str) and self.init_weights.lower() in {'true', 'false'}:
+            self.init_weights = bool(strtobool(self.init_weights))
+        self._init_multimodal_full()
+        if self.target_regex:
+            self.target_modules = self.target_regex
+
+    def _init_multimodal_full(self):
+        model_arch = self.model_meta.model_arch
+        if not self.model_meta.is_multimodal or not model_arch or self.tuner_type != 'full':
+            return
+        if self.freeze_llm:
+            self.freeze_parameters += model_arch.language_model
+        if self.freeze_vit:
+            self.freeze_parameters += model_arch.vision_tower
+        if self.freeze_aligner:
+            self.freeze_parameters += model_arch.aligner
+        else:
+            self.trainable_parameters += model_arch.aligner
+        self.freeze_parameters += model_arch.generator
+        if self.freeze_parameters:
+            logger.info(f'freeze_parameters: {self.freeze_parameters}')
+        if self.trainable_parameters:
+            logger.info(f'additional trainable_parameters: {self.trainable_parameters}')

swift/arguments/webui_args.py

@@ -0,0 +1,18 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from dataclasses import dataclass
+
+
+@dataclass
+class WebUIArguments:
+    """A dataclass for web UI configuration arguments.
+
+    Args:
+        server_name (str): The hostname or IP address to be bound to the Web UI server. Defaults to '0.0.0.0'.
+        server_port (int): The port number to be bound to the Web UI server. Defaults to 7860.
+        share (bool): Whether to create a public, shareable link for the web UI. Defaults to False.
+        lang (str): The language for the web UI, chosen from {'zh', 'en'}. Defaults to 'zh'.
+    """
+    server_name: str = '0.0.0.0'
+    server_port: int = 7860
+    share: bool = False
+    lang: str = 'zh'

swift/callbacks/__init__.py

@@ -0,0 +1,2 @@
+from .base import TrainerCallback
+from .mapping import callbacks_map

swift/callbacks/activation_cpu_offload.py

@@ -0,0 +1,607 @@
+"""Functionality for CPU offloading of tensors saved for backward pass."""
+import functools
+import torch
+from torch.distributed.fsdp import FSDPModule as FSDP2
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from transformers.trainer_callback import TrainerControl, TrainerState
+from transformers.training_args import TrainingArguments
+from transformers.utils import is_torch_npu_available
+from typing import Any, Optional
+
+from swift.utils import get_logger
+from .base import TrainerCallback
+
+logger = get_logger()
+
+is_cuda_available = torch.cuda.is_available()
+is_npu_available = is_torch_npu_available()
+
+
+def _get_unique_tensor_key(tensor):
+    key = (tensor.untyped_storage().data_ptr() + tensor.storage_offset(), tensor.dtype)
+    return key
+
+
+def get_device_name() -> str:
+    """Function that gets the torch.device based on the current machine.
+    This currently only supports CPU, CUDA, NPU.
+    Returns:
+        device
+    """
+    if is_cuda_available:
+        device = 'cuda'
+    elif is_npu_available:
+        device = 'npu'
+    else:
+        device = 'cpu'
+    return device
+
+
+class FSDPParameterFilter:
+
+    def __init__(self):
+        self.model_parameters_storage = set()
+
+    def __call__(self, tensor):
+        return tensor.untyped_storage().data_ptr() not in self.model_parameters_storage
+
+    def update_model_parameters(self, model):
+        new_storage = set()
+        for p in model.parameters():
+            new_storage.add(p.data.untyped_storage().data_ptr())
+        self.model_parameters_storage = new_storage
+
+
+def get_torch_device() -> Any:
+    """Return the corresponding torch attribute based on the device type string.
+    Returns:
+        module: The corresponding torch device namespace, or torch.cuda if not found.
+    """
+    device_name = get_device_name()
+    try:
+        return getattr(torch, device_name)
+    except AttributeError:
+        logger.warning(f"Device namespace '{device_name}' not found in torch, try to load torch.cuda.")
+        return torch.cuda
+
+
+class CpuOffloadHookWithOffloadHandler:
+    """Context-manager that offloads/recovers tensors through an offload hander.
+
+    The hook just offloads/recovers the tensor object to the handler through `tensor_push`
+    and `tensor_pop` interface. How the offload-handler manages the offloading, recovering
+    or prefetching timing is transparent to this hook.
+    """
+
+    def __init__(
+        self,
+        offload_handler: 'OffloadHandler',
+        handler_extra_kwargs: Optional[dict[str, Any]] = None,
+    ) -> None:
+        if handler_extra_kwargs is None:
+            handler_extra_kwargs = {}
+        self.offload_handler: OffloadHandler = offload_handler
+        self.handler_extra_kwargs: dict[str, Any] = handler_extra_kwargs
+        self.inside_context = False
+
+    def __enter__(self):
+        self.inside_context = True
+        torch._C._autograd._push_saved_tensors_default_hooks(self.on_save_for_backward, self.on_get_saved_tensor)
+
+    def __exit__(self, *args: Any):
+        self.inside_context = False
+        torch._C._autograd._pop_saved_tensors_default_hooks()
+
+    def on_save_for_backward(self, tensor: torch.Tensor) -> Any:
+        retrieve_identifier = self.offload_handler.tensor_push(tensor, **self.handler_extra_kwargs)
+        return retrieve_identifier
+
+    def on_get_saved_tensor(self, saved_state: Any) -> torch.Tensor:
+        tensor = self.offload_handler.tensor_pop(saved_state, **self.handler_extra_kwargs)
+        return tensor
+
+
+class OffloadHandler:
+    """A base class for CPU offload-handler."""
+
+    def __init__(self) -> None:
+        pass
+
+    def tensor_push(self, tensor: torch.Tensor, **kwargs) -> Any:
+        """Tensor push."""
+        raise NotImplementedError(
+            '`tensor_push is not implented in OffloadHandler class. Inherit this class and implement your '
+            'custom tensor_push.')
+
+    def tensor_pop(self, tensor_tag: Any, **kwargs):
+        """Tensor pop."""
+        raise NotImplementedError(
+            '`tensor_pop is not implented in OffloadHandler class. Inherit this class and implement your '
+            'custom tensor_pop.')
+
+
+class GroupCommitFunction(torch.autograd.Function):
+    """this is a dummy op with output identical to input.
+    However, it is necessary for marking a timepoint for offload handler to
+    accomplish all synchronizations. Implementing it as a function is necessary
+    because we need to actions in both forward and backward.
+    """
+
+    @staticmethod
+    def forward(ctx, tensor, cpu_offload_handler):
+        # pylint: disable=missing-function-docstring
+        cpu_offload_handler.on_group_commit_forward()
+        ctx.cpu_offload_handler = cpu_offload_handler
+        # return the identical tensor
+        return tensor
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        # pylint: disable=missing-function-docstring
+        cpu_offload_handler = ctx.cpu_offload_handler
+        cpu_offload_handler.on_group_commit_backward()
+        return grad_output, None
+
+
+group_prefetch_offload_commit = GroupCommitFunction.apply
+
+
+class SynchronizedGroupOffloadHandler(OffloadHandler):
+    """Offload Handler that offloads/reloads in a synchronized way.
+    The device-to-host and host-to-device copying happen in the same stream
+    as the computation kernels, thus the copying will block computation.
+    """
+
+    def __init__(self, num_offload_group, tensor_need_offloading_checker=(lambda _: True)) -> None:
+        super().__init__()
+
+        self.num_offload_group = num_offload_group
+        self.tensor_need_offloading_checker = tensor_need_offloading_checker
+
+        self.groupid_reset()
+
+    def groupid_reset(self):
+        """Groupid reset."""
+        # Data structures to label saved tensors and book-keep their cpu copies.
+        # Currently, on push, create a new cpu tensor and copies; on pop, copies
+        # the tensor back to gpu and deletes the cpu tensor.
+        # These will increment whenever `group_commit()` is invoked
+        self.current_group, self.tensor_count_current_group = (0, 0)
+        self.torch_tensor_count = 0
+        self.tensor_tag_to_state = {}
+
+    def on_group_commit_forward(self):
+        """On group commit forward."""
+        # finishing up with updating current group and tensor count
+        self.current_group += 1  # increment
+        self.tensor_count_current_group = 0  # reset
+
+    def on_group_commit_backward(self):
+        """On group commit backward."""
+        self.current_group -= 1
+        assert self.current_group >= 0
+
+    @staticmethod
+    def offload(src_tensor, pin_memory=True):
+        """Offload."""
+        # NPU doesn't fully support async H2D/D2H with pinned memory; use sync copy.
+        if is_npu_available:
+            cpu_backup = torch.empty(
+                src_tensor.size(),
+                dtype=src_tensor.dtype,
+                layout=src_tensor.layout,
+                device='cpu',
+                pin_memory=False,
+            )
+            cpu_backup.copy_(src_tensor, non_blocking=False)
+        else:
+            cpu_backup = torch.empty(
+                src_tensor.size(),
+                dtype=src_tensor.dtype,
+                layout=src_tensor.layout,
+                device='cpu',
+                pin_memory=pin_memory,
+            )
+            cpu_backup.copy_(src_tensor, non_blocking=True)
+        state = (src_tensor.device, cpu_backup)
+        return state
+
+    @staticmethod
+    def reload(state, non_blocking=None):
+        """Reload."""
+        dev, cpu_backup = state
+        if non_blocking is None:
+            non_blocking = cpu_backup.is_pinned()
+        return cpu_backup.to(dev, non_blocking=non_blocking)
+
+    def tensor_push(self, tensor: torch.Tensor, **kwargs):
+        """Tensor push."""
+        # obtain a unique tensor tag
+        tensor_tag = (self.current_group, self.tensor_count_current_group)
+        self.tensor_count_current_group += 1
+        assert tensor_tag not in self.tensor_tag_to_state
+        if self.current_group < self.num_offload_group and self.tensor_need_offloading_checker(tensor):
+            state = SynchronizedGroupOffloadHandler.offload(tensor)
+            self.tensor_tag_to_state[tensor_tag] = state
+        else:
+            # will be offloaded together after group commit
+            self.tensor_tag_to_state[tensor_tag] = tensor
+
+        return tensor_tag
+
+    def tensor_pop(self, tensor_tag, **kwargs):
+        """Tensor pop."""
+        assert tensor_tag in self.tensor_tag_to_state
+        state = self.tensor_tag_to_state.pop(tensor_tag)
+        if isinstance(state, tuple):
+            tensor = SynchronizedGroupOffloadHandler.reload(state)
+        else:
+            tensor = state
+        return tensor
+
+
+class AsyncDoubleBufferGroupOffloadHandler(SynchronizedGroupOffloadHandler):
+    """Compared to synchronize, this uses more memory because of the buffer but
+    achieves better performance due to the overlapping. D2h and h2d copying are
+    completely hidden behind computation if computation time of a layer is longer
+    than host-device communication time. Bulk offloading with delay and bulk reloading
+    with prefetch are implemented."""
+
+    def __init__(
+            self,
+            num_offload_group,  # must be <= actual number of groups (number of commits)
+            num_model_group,
+            tensor_need_offloading_checker=(lambda t: True),
+    ) -> None:
+        super().__init__(
+            num_offload_group=num_offload_group,
+            tensor_need_offloading_checker=tensor_need_offloading_checker,
+        )
+        # Number of layers in the model
+        self.num_layers = num_model_group
+        # Data Structure to maintain reference to activation tensors
+        self.tensor_tag_to_buf = {}
+        # Tracking the number of layers offloaded
+        self.offloaded_group_count = 0
+        # Core data structure that decides the window for offloading
+        self.layer_window_map = {}
+        self.group_offload_mapping = {}
+
+        # Logic to make offloading load balance across computation
+        # for optimal CPU/GPU interconnect usage
+        constant = 0
+        for i in range(self.num_offload_group):
+            self.layer_window_map[i] = ((self.num_layers // self.num_offload_group) * (i + 1)) - 1
+            if i < (self.num_layers % self.num_offload_group):
+                self.layer_window_map[i] += i + 1
+                constant = i + 1
+            else:
+                self.layer_window_map[i] += constant
+
+        # allocate streams and events for synchronization
+        self.d2h_stream = get_torch_device().Stream()
+        self.h2d_stream = get_torch_device().Stream()
+
+    def tensor_push(self, tensor: torch.Tensor, **kwargs) -> Any:
+        torch_stray_tensor = isinstance(
+            tensor,
+            torch._subclasses.fake_tensor.FakeTensor | torch._subclasses.functional_tensor.FunctionalTensor,
+        )
+        need_offload = not torch_stray_tensor
+        need_offload = need_offload and self.tensor_need_offloading_checker(tensor)
+
+        if need_offload:
+            # obtain a unique tensor tag
+            tensor_tag = (self.current_group, self.tensor_count_current_group)
+            self.tensor_count_current_group += 1
+
+            assert tensor_tag not in self.tensor_tag_to_state
+            self.tensor_tag_to_state[tensor_tag] = tensor
+
+            if self.current_group < self.num_offload_group:
+                self.tensor_tag_to_buf[tensor_tag] = tensor
+        else:
+            tensor_tag = tensor
+        return tensor_tag
+
+    def tensor_pop(self, tensor_tag, **kwargs):
+        """Tensor pop."""
+        if isinstance(tensor_tag, torch.Tensor):
+            return tensor_tag
+        assert tensor_tag in self.tensor_tag_to_state
+        tensor = self.tensor_tag_to_state.pop(tensor_tag)
+        self.tensor_tag_to_buf.pop(tensor_tag, None)
+
+        # the tensor should have been copied back in on_group_commit_backward()
+        # which invokes bulk_reload_group.
+        assert not isinstance(tensor, tuple)
+        return tensor
+
+    def bulk_offload_group(self, group_to_offload):
+        """Bulk offload group."""
+        offload_mapping = {}
+        offload_size = 0
+        with get_torch_device().stream(self.d2h_stream):
+            for tensor_tag, state in self.tensor_tag_to_state.items():
+                group_id, _ = tensor_tag
+                if group_id == group_to_offload:
+                    assert not isinstance(state, tuple)
+                    key = _get_unique_tensor_key(state)
+                    if key not in offload_mapping:
+                        offload_mapping[key] = state
+                    # if offload, return the reference to cpu copy
+                    self.tensor_tag_to_state[tensor_tag] = (key, state.shape)
+            for key, tensor in offload_mapping.items():
+                state = SynchronizedGroupOffloadHandler.offload(tensor)
+                offload_size += tensor.numel() * tensor.element_size()
+                offload_mapping[key] = state
+
+            self.group_offload_mapping[group_to_offload] = offload_mapping
+
+    def synchronize_on_group_commit_forward(self, current_group):
+        """Synchronize on group commit forward."""
+
+        # For the first group, kickstart the offload after we have
+        # the first compute completion
+        if current_group == 0:
+            self.d2h_stream.wait_stream(get_torch_device().current_stream())
+            self.bulk_offload_group(current_group)
+
+        # Window map data structure helps us synchronize based on number
+        # of layers offloaded
+        if self.layer_window_map[self.offloaded_group_count] == current_group:
+            # Stream synchronization both ways
+            self.d2h_stream.wait_stream(get_torch_device().current_stream())
+            get_torch_device().current_stream().wait_stream(self.d2h_stream)
+
+            # Time to free the activation memory after usage
+            for tensor_tag, _ in self.tensor_tag_to_buf.items():
+                if tensor_tag[0] == self.offloaded_group_count:
+                    self.tensor_tag_to_buf[tensor_tag] = None
+
+            # Time to offload the next group
+            if self.offloaded_group_count < (self.num_offload_group - 1):
+                self.bulk_offload_group(self.offloaded_group_count + 1)
+
+            # Increment the offload group count to keep track
+            self.offloaded_group_count += 1
+
+    def on_group_commit_forward(self):
+        """This function will cause host device synchronization"""
+        # handle synchronization events
+        self.synchronize_on_group_commit_forward(self.current_group)
+
+        super().on_group_commit_forward()
+
+    @torch.no_grad
+    def bulk_reload_group(self, group_to_reload):
+        """Bulk reload group."""
+        assert group_to_reload < self.num_offload_group
+
+        with get_torch_device().stream(self.h2d_stream):
+            # move back tensors
+            offload_mapping = self.group_offload_mapping.pop(group_to_reload)
+            assert offload_mapping is not None
+            for key, state in offload_mapping.items():
+                offload_mapping[key] = SynchronizedGroupOffloadHandler.reload(state)
+            for tensor_label, state in self.tensor_tag_to_state.items():
+                group_id, _ = tensor_label
+                if group_id == group_to_reload and not isinstance(state, torch.Tensor):
+                    assert isinstance(state, tuple), f'{group_id} {state}'
+                    key, shape = state
+                    recovered_tensor = offload_mapping[key].view(shape)
+                    self.tensor_tag_to_state[tensor_label] = recovered_tensor
+
+    def on_group_commit_backward(self):
+        # first decrement the current group.
+        # after last commit in forward, the group will +1; in backward it -1.
+        # Finally it should be decremented to 0.
+        self.current_group -= 1
+        assert self.current_group >= 0
+
+        # Layer window data structure helps us to reload at right times
+        if self.layer_window_map[self.offloaded_group_count - 1] == self.current_group:
+            # Stream synchronization both ways
+            self.h2d_stream.wait_stream(get_torch_device().current_stream())
+            get_torch_device().current_stream().wait_stream(self.h2d_stream)
+
+            # Time to reload the next group
+            self.bulk_reload_group(self.offloaded_group_count - 1)
+
+            # Decrease the offloading group counter
+            self.offloaded_group_count -= 1 if self.offloaded_group_count > 1 else 0
+
+        # Last group computation needs to wait till all the reloads complete
+        if self.current_group == 0:
+            get_torch_device().current_stream().wait_stream(self.h2d_stream)
+            self.offloaded_group_count = 0
+
+
+def get_activation_offload_context(num_layers: int = 1,
+                                   model_layers: int = 1,
+                                   tensor_need_offloading_checker=(lambda t: True)):
+    cpu_offload_handler = AsyncDoubleBufferGroupOffloadHandler(
+        num_offload_group=num_layers,
+        num_model_group=model_layers,
+        tensor_need_offloading_checker=tensor_need_offloading_checker,
+    )
+
+    def group_prefetch_offload_commit_async(tensor):
+        return group_prefetch_offload_commit(tensor, cpu_offload_handler)
+
+    return (
+        CpuOffloadHookWithOffloadHandler(offload_handler=cpu_offload_handler),
+        group_prefetch_offload_commit_async,
+    )
+
+
+class ActivationHandler:
+
+    def __init__(self, offload_ctx, sync_func, tensor_filter, enable_ckpt):
+        self._offload_ctx = offload_ctx
+        self._sync_func = sync_func
+        self._enable_ckpt = enable_ckpt
+        self._tensor_filter = tensor_filter
+        if enable_ckpt:
+            self.checkpoint_fn = functools.partial(
+                torch.utils.checkpoint.checkpoint,
+                use_reentrant=True,
+            )
+
+    def pre_forward(self, module):
+        if module.training:
+            self._offload_ctx.__enter__()
+            self._tensor_filter.update_model_parameters(module)
+
+    def post_forward(self, module):
+        if module.training:
+            self._offload_ctx.__exit__(None, None, None)
+
+    def _pack_kwargs(self, *args, **kwargs):
+        kwarg_keys = []
+        flat_args = list(args)
+        for k, v in kwargs.items():
+            kwarg_keys.append(k)
+            flat_args.append(v)
+
+        return tuple(flat_args), tuple(kwarg_keys)
+
+    def _unpack_kwargs(self, flat_args, kwarg_keys):
+        assert len(kwarg_keys) <= len(flat_args), f'too many keys {len(kwarg_keys)} vs. {len(flat_args)}'
+        if len(kwarg_keys) == 0:
+            return flat_args, {}
+        args = flat_args[:-len(kwarg_keys)]
+        kwargs = dict(zip(kwarg_keys, flat_args[-len(kwarg_keys):], strict=True))
+        return args, kwargs
+
+    def _ckpt_forward(self, forward_method, *args, **kwargs):
+        flat_args, kwarg_keys = self._pack_kwargs(*args, **kwargs)
+
+        def my_function(*inputs):
+            # unpack back into args and kwargs
+            unpacked_args, unpacked_kwargs = self._unpack_kwargs(inputs, kwarg_keys)
+            # run original module
+            return forward_method(*unpacked_args, **unpacked_kwargs)
+
+        return self.checkpoint_fn(
+            my_function,
+            *flat_args,
+        )
+
+    def forward(self, module, forward_method, *args, **kwargs):
+        if not module.training:
+            return forward_method(*args, **kwargs)
+        if not self._enable_ckpt:
+            ret = forward_method(*args, **kwargs)
+        else:
+            ret = self._ckpt_forward(forward_method, *args, **kwargs)
+        binded_tensor = ret
+        if isinstance(ret, tuple):
+            binded_tensor = ret[0]
+        binded_tensor = self._sync_func(binded_tensor)
+        final_ret = binded_tensor
+        if isinstance(ret, tuple):
+            final_ret = (final_ret, ) + ret[1:]
+        return final_ret
+
+    def wrap_module_forward_method(self, module):
+        orig_method = module.forward
+        handler = self
+
+        @functools.wraps(orig_method)
+        def wrapped_method(model_self, *args, **kwargs):
+            handler.pre_forward(model_self)
+            out = handler.forward(model_self, orig_method, *args, **kwargs)
+            handler.post_forward(model_self)
+            return out
+
+        module.forward = wrapped_method.__get__(module, type(module))
+
+
+def enable_activation_offloading(model, strategy, enable_ckpt=False):
+    """
+    Enable activation offloading for the model. It groups activations by TransformerLayer and offloads activation
+    groups asynchronously. This means that the offloading of the i-th activation group and the computation of the i+1-th
+    activation group happen at the same time, and there are at most two activation groups in GPU memory.
+
+    Args:
+        model: the model to enable activation offloading
+        strategy: the training strategy of the model, such as "fsdp"
+        enable_ckpt: whether activation checkpointing(also called gradient checkpointing) has been enabled for the model
+
+    Note:
+        For best efficiency, activation offloading is usually combined with activation checkpointing. However, this
+        implementation of activation offloading is conflicted with the implementation of activation checkpointing in
+        some training strategies. This function resolves this conflict, and therefore requires the "strategy" and
+        "enable_ckpt" arguments.
+
+    Returns:
+
+    """
+
+    assert strategy == 'fsdp' or strategy == 'fsdp2', 'activation offloading only supports fsdp strategy'
+    layers = []
+
+    def get_layers(module):
+        for name, child in module.named_children():
+            if not isinstance(child, FSDP | FSDP2):
+                get_layers(child)
+            else:
+                wrapped_module = child
+                if isinstance(child, FSDP):
+                    wrapped_module = child._fsdp_wrapped_module
+                # In some cases, torch.nn.Embedding is wrapped with FSDP alone. However, the activation
+                # size of torch.nn.Embedding is small, so it's not necessary to offload it.
+                if not isinstance(wrapped_module, torch.nn.Embedding):
+                    layers.append(child)
+
+    get_layers(model)
+    if len(layers) < 3:
+        logger.warning(f'Find only {len(layers)} fsdp layers, not necessary to enable async activation offloading')
+        return
+
+    tensor_filter = FSDPParameterFilter()
+    context, sync_func = get_activation_offload_context(len(layers) - 1, len(layers), tensor_filter)
+    if enable_ckpt:
+        # The implementation of activation checkpointing in transformers library is incompatible with
+        # activation offloading,
+        # so it will be disabled, but this implementation supports another version of activation checkpointing, so that
+        # these two features can be enabled at the same time.
+        for module in model.modules():
+            if hasattr(module, 'gradient_checkpointing_disable'):
+                module.gradient_checkpointing_disable()
+
+    handler = ActivationHandler(context, sync_func, tensor_filter, enable_ckpt)
+    for layer in layers:
+        module = layer
+        if isinstance(layer, FSDP):
+            module = module._fsdp_wrapped_module
+        handler.wrap_module_forward_method(module)
+
+
+class ActivationCpuOffloadCallBack(TrainerCallback):
+
+    def __init__(self, args: TrainingArguments, trainer):
+        super().__init__(args, trainer)
+
+    def on_train_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
+        """
+        Event called at the beginning of training.
+        """
+        model = kwargs['model']
+
+        # Check if model is wrapped with FSDP
+        if isinstance(model, FSDP) or isinstance(model, FSDP2):
+            if args is not None and hasattr(args, 'fsdp_config'):
+                fsdp_config = args.fsdp_config
+                # Check if fsdp_config is a dictionary and has activation_cpu_offload enabled
+                if isinstance(fsdp_config, dict) and fsdp_config.get('activation_cpu_offload', False):
+                    # Get FSDP version from fsdp_config
+                    strategy = fsdp_config.get('version', None)
+                    if strategy is not None:
+                        fsdp_version = 'fsdp' if strategy == 1 else 'fsdp2'
+                        # Get activation checkpointing setting from fsdp_config
+                        enable_ckpt = fsdp_config.get('activation_checkpointing', False)
+                        if enable_ckpt and hasattr(model, 'enable_input_require_grads'):
+                            model.enable_input_require_grads()
+                        enable_activation_offloading(model, strategy=fsdp_version, enable_ckpt=enable_ckpt)

swift/callbacks/adalora.py

@@ -0,0 +1,31 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import types
+from typing import TYPE_CHECKING
+
+from .base import TrainerCallback
+
+if TYPE_CHECKING:
+    from swift.trainers import Trainer, TrainingArguments
+
+
+class AdaloraCallback(TrainerCallback):
+
+    def __init__(self, args: 'TrainingArguments', trainer: 'Trainer'):
+        super().__init__(args, trainer)
+        self.global_step = 0
+        self.args = args
+
+    # offload original_modules to cpu, to save memory
+    def on_train_begin(self, _args, state, control, **kwargs):
+        model = kwargs['model']
+        model.peft_config['default'].total_step = state.max_steps
+
+        def zero_grad(_self, *args, **kwargs):
+            _self.update_and_allocate(self.global_step + 1)
+            _self._zero_grad(*args, **kwargs)
+
+        model._zero_grad = model.zero_grad
+        model.zero_grad = types.MethodType(zero_grad, model)
+
+    def on_step_end(self, _args, state, control, **kwargs):
+        self.global_step = state.global_step

swift/callbacks/base.py

@@ -0,0 +1,13 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from transformers import TrainerCallback as HfTrainerCallback
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from swift.trainers import Trainer, TrainingArguments
+
+
+class TrainerCallback(HfTrainerCallback):
+
+    def __init__(self, args: 'TrainingArguments', trainer: 'Trainer'):
+        self.args = args
+        self.trainer = trainer

swift/callbacks/deepspeed_elastic.py

@@ -0,0 +1,45 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import torch
+import torch.distributed as dist
+
+from swift.utils import ShutdownManager, get_device
+from .base import TrainerCallback
+
+
+class DeepspeedElasticCallback(TrainerCallback):
+    """Compatibility marker for enabling DeepSpeed elastic setup during argument initialization."""
+
+
+class GracefulExitCallback(TrainerCallback):
+
+    def __init__(self, args=None, trainer=None):
+        if args is not None and trainer is not None:
+            super().__init__(args, trainer)
+        shutdown_manager = ShutdownManager()
+        shutdown_manager.register()
+        self.shutdown_manager = shutdown_manager
+        self._pending_stop = False
+
+    def on_step_end(self, args, state, control, **kwargs):
+        device_type = get_device()
+
+        local_req = 1 if self.shutdown_manager.should_shutdown() else 0
+        if dist.is_available() and dist.is_initialized():
+
+            t = torch.tensor([local_req], dtype=torch.uint8, device=device_type)
+            # all_reduce with MAX: if any rank has 1 -> result 1 everywhere
+            dist.all_reduce(t, op=dist.ReduceOp.MAX)
+            any_req = bool(int(t.item()))
+        else:
+            any_req = bool(local_req)
+
+        if any_req:
+            control.should_save = True
+            self._pending_stop = True
+        return control
+
+    def on_save(self, args, state, control, **kwargs):
+        if self._pending_stop:
+            control.should_training_stop = True
+            self._pending_stop = False
+        return control

swift/callbacks/early_stop.py

@@ -0,0 +1,34 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import numpy as np
+from transformers import TrainerControl, TrainerState
+from typing import TYPE_CHECKING
+
+from swift.utils import get_logger
+from .base import TrainerCallback
+
+if TYPE_CHECKING:
+    from swift.trainers import Trainer, TrainingArguments
+
+logger = get_logger()
+
+
+class EarlyStopCallback(TrainerCallback):
+    """An early stop implementation"""
+
+    def __init__(self, args: 'TrainingArguments', trainer: 'Trainer'):
+        super().__init__(args, trainer)
+        self.best_metric = None
+        self.interval = 0
+        self.total_interval = args.early_stop_interval
+
+    def on_save(self, args: 'TrainingArguments', state: TrainerState, control: TrainerControl, **kwargs):
+        operator = np.greater if args.greater_is_better else np.less
+        if self.best_metric is None or operator(state.best_metric, self.best_metric):
+            self.best_metric = state.best_metric
+            self.interval = 0
+        else:
+            self.interval += 1
+
+        if self.interval >= self.total_interval:
+            logger.info(f'Training stop because of eval metric is stable at step {state.global_step}')
+            control.should_training_stop = True

swift/callbacks/lisa.py

@@ -0,0 +1,57 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import numpy as np
+import torch
+from typing import TYPE_CHECKING
+
+from .base import TrainerCallback
+
+if TYPE_CHECKING:
+    from swift.trainers import Trainer, TrainingArguments
+
+
+class LISACallback(TrainerCallback):
+
+    def __init__(self, args: 'TrainingArguments', trainer: 'Trainer'):
+        assert args.tuner_type == 'full', 'LISA only supports full parameter training.'
+        super().__init__(args, trainer)
+        self.n_layers = args.lisa_activated_layers
+        self.step_interval = args.lisa_step_interval
+        self.model = self.trainer.model
+        layers_name = None
+        layers = None
+        for name, module in self.model.named_modules():
+            if isinstance(module, torch.nn.ModuleList):
+                layers_name = name
+                layers = module
+                break
+        assert layers_name is not None
+        self.layers_attribute = layers_name
+        self.total_layers = len(layers)
+
+        # Freeze all layers upon initialization
+        self.freeze_all_layers()
+        self.active_layers_indices = []
+        self.switch_active_layers()
+
+    def freeze_all_layers(self):
+        layers = self.model.get_submodule(self.layers_attribute)
+        for layer in layers:
+            for param in layer.parameters():
+                param.requires_grad = False
+
+    def on_step_begin(self, args, state, control, **kwargs):
+        # Check if it's time to switch active layers, including at step 0
+        if state.global_step % self.step_interval == 0 or state.global_step == 1:
+            self.switch_active_layers()
+
+    def switch_active_layers(self):
+        # First, disable gradients for all layers
+        self.freeze_all_layers()
+
+        # Randomly select n_layers to activate
+        layers = self.model.get_submodule(self.layers_attribute)
+        self.active_layers_indices = np.random.choice(range(self.total_layers), self.n_layers, replace=False)
+        # Enable gradients only for the selected layers
+        for idx in self.active_layers_indices:
+            for param in layers[idx].parameters():
+                param.requires_grad = True

swift/callbacks/mapping.py

@@ -0,0 +1,17 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from .activation_cpu_offload import ActivationCpuOffloadCallBack
+from .adalora import AdaloraCallback
+from .deepspeed_elastic import DeepspeedElasticCallback, GracefulExitCallback
+from .early_stop import EarlyStopCallback
+from .lisa import LISACallback
+from .perf_log import PerfMetricsLogCallback
+
+callbacks_map = {
+    'activation_cpu_offload': ActivationCpuOffloadCallBack,
+    'adalora': AdaloraCallback,
+    'deepspeed_elastic': DeepspeedElasticCallback,
+    'early_stop': EarlyStopCallback,
+    'graceful_exit': GracefulExitCallback,
+    'lisa': LISACallback,
+    'perf_log': PerfMetricsLogCallback
+}

swift/callbacks/perf_log.py

@@ -0,0 +1,137 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import time
+import torch
+from transformers import TrainerControl, TrainerState
+from typing import TYPE_CHECKING
+
+from swift.utils import (empty_cache, get_current_device, get_device_count, get_dist_setting, get_env_args, get_logger,
+                         synchronize)
+from .base import TrainerCallback
+
+if TYPE_CHECKING:
+    from swift.trainers import Trainer, TrainingArguments
+
+logger = get_logger()
+
+device_flops_map = {
+    'GB200': 2.5e15,
+    'B200': 2.25e15,
+    'MI300X': 1336e12,
+    'H100': 312e12,
+    'H800': 312e12,
+    'H200': 989e12,
+    'A100': 312e12,
+    'A800': 312e12,
+    'L40S': 362.05e12,
+    'L40': 181.05e12,
+    'A40': 149.7e12,
+    'L20': 119.5e12,
+    'H20': 148e12,
+    '910B': 354e12,
+    'Ascend910': 354e12,
+    'RTX 3070 Ti': 21.75e12
+}
+
+
+class PerfMetricsLogCallback(TrainerCallback):
+    """An callback for perf metrics (MFU etc) log implementation"""
+
+    def __init__(self, args: 'TrainingArguments', trainer: 'Trainer'):
+        super().__init__(args, trainer)
+        self.max_tflops = None
+        self.elapsed = 0.0
+        self.step_start_time = None
+
+    def on_init_end(self, args: 'TrainingArguments', state: TrainerState, control: TrainerControl, **kwargs):
+
+        # Top priority. Specify by ENV
+        tflops = get_env_args('DEVICE_TFLOPS', float, None)
+        # `state.total_flos` is summed across all ranks (cluster-global) by HF Trainer,
+        # so the theoretical denominator must use the TOTAL number of GPUs in use across the entire cluster.
+        _, _, world_size, local_world_size = get_dist_setting()
+        local_n_gpu = get_device_count()
+        gpus_per_process = max(1, local_n_gpu // max(local_world_size, 1))
+        device_count = max(world_size * gpus_per_process, 1)
+        if tflops is not None:
+            logger.info(f"Specify theoretical max TFLOPS through ENV 'DEVICE_TFLOPS'. [{tflops} TFLOPS]")
+        else:
+            # Run a estimating test.
+            dtype = kwargs.get('model').dtype
+            device = torch.device(get_current_device())
+            logger.info(f'Estimating device TFLOPS baseline. Device: [{device}] dtype: [{dtype}]')
+            tflops = self._estimate_device_tflops_by_dtype(device, dtype)
+            logger.info(f'Estimate test finished. [{tflops} TFLOPS] Device count: [{device_count}]')
+
+        self.max_tflops = tflops * device_count
+
+    def on_step_begin(self, args: 'TrainingArguments', state: TrainerState, control: TrainerControl, **kwargs):
+        self.step_start_time = time.time()
+
+    def on_step_end(self, args: 'TrainingArguments', state: TrainerState, control: TrainerControl, **kwargs):
+        self.elapsed += time.time() - self.step_start_time
+
+    def on_log(self, args: 'TrainingArguments', state: TrainerState, control: TrainerControl, logs=None, **kwargs):
+        total_flos = getattr(state, 'total_flos', 0)
+        actual_flops = total_flos / self.elapsed
+        theoretical_max_flops = self.max_tflops * 1e12
+        mfu = actual_flops / theoretical_max_flops
+        logger.debug(f'Total_flos[{total_flos}] elapsed_time[{self.elapsed}]sec Average MFU[{mfu}]')
+        logs['MFU'] = round(mfu, 6)
+
+    @staticmethod
+    def _estimate_device_tflops_by_dtype(device: torch.device, dtype: torch.dtype, repeats: int = 60, dim: int = 8192):
+        # Set matrix dimension
+        shape = (dim, dim)
+        backend = device.type
+        if backend == 'npu':
+            import torch_npu
+
+        # Initialize matrices
+        a = torch.randn(*shape, device=device, dtype=dtype)
+        b = torch.randn(*shape, device=device, dtype=dtype)
+
+        # Warm-up
+        for _ in range(5):
+            c = torch.matmul(a, b)
+        synchronize(device)
+
+        # Run benchmark test
+        start = time.time()
+        for _ in range(repeats):
+            c = torch.matmul(a, b)
+        synchronize(device)
+        end = time.time()
+        total_time = end - start
+        avg_time = total_time / repeats
+
+        # Adjust repeat count and retest if test duration is too short
+        if total_time < 3:
+            repeats = int(6 / avg_time)
+            start = time.time()
+            for _ in range(repeats):
+                c = torch.matmul(a, b)
+            synchronize(device)
+            end = time.time()
+            total_time = end - start
+            avg_time = total_time / repeats
+
+        del a, b, c
+        empty_cache()
+
+        tflops = (2 * dim**3 / avg_time) / 1e12
+        logger.info(f'[Device {device}] Total time: {total_time:.4f}s, dtype: {dtype}, Perf: {tflops:.4f} TFLOPS')
+        return tflops
+
+    @staticmethod
+    def _retrieve_flops_from_map(device):
+        """Retrieve theoretical FLOPS from Map.    """
+        # This function is never used.
+
+        device_name = device.get_device_name()
+        flops = None
+        for name, value in device_flops_map.items():
+            if name in device_name:
+                flops = value
+                break
+
+        return flops

swift/cli/_megatron/export.py

@@ -0,0 +1,7 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import os
+
+if __name__ == '__main__':
+    os.environ.setdefault('CUDA_DEVICE_MAX_CONNECTIONS', '1')
+    from swift.megatron import megatron_export_main
+    megatron_export_main()

swift/cli/_megatron/main.py

@@ -0,0 +1,19 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from typing import Dict
+
+from ..main import cli_main as swift_cli_main
+
+ROUTE_MAPPING: Dict[str, str] = {
+    'pt': 'swift.cli._megatron.pt',
+    'sft': 'swift.cli._megatron.sft',
+    'rlhf': 'swift.cli._megatron.rlhf',
+    'export': 'swift.cli._megatron.export',
+}
+
+
+def cli_main():
+    return swift_cli_main(ROUTE_MAPPING, is_megatron=True)
+
+
+if __name__ == '__main__':
+    cli_main()

swift/cli/_megatron/pt.py

@@ -0,0 +1,7 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import os
+
+if __name__ == '__main__':
+    os.environ.setdefault('CUDA_DEVICE_MAX_CONNECTIONS', '1')
+    from swift.megatron import megatron_pretrain_main
+    megatron_pretrain_main()

swift/cli/_megatron/rlhf.py

@@ -0,0 +1,24 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import os
+import sys
+
+
+def _use_ray() -> bool:
+    if '--use_ray' not in sys.argv:
+        return False
+    idx = sys.argv.index('--use_ray')
+    sys.argv.pop(idx)
+    if idx < len(sys.argv) and sys.argv[idx].lower() in ('true', 'false'):
+        val = sys.argv.pop(idx).lower() == 'true'
+        return val
+    return True
+
+
+if __name__ == '__main__':
+    if _use_ray():
+        from swift.ray.megatron.pipeline import main as ray_main
+        ray_main()
+    else:
+        os.environ.setdefault('CUDA_DEVICE_MAX_CONNECTIONS', '1')
+        from swift.megatron import megatron_rlhf_main
+        megatron_rlhf_main()

swift/cli/_megatron/sft.py

@@ -0,0 +1,7 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import os
+
+if __name__ == '__main__':
+    os.environ.setdefault('CUDA_DEVICE_MAX_CONNECTIONS', '1')
+    from swift.megatron import megatron_sft_main
+    megatron_sft_main()

swift/cli/app.py

@@ -0,0 +1,5 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from swift.pipelines import app_main
+
+if __name__ == '__main__':
+    app_main()

swift/cli/deploy.py

@@ -0,0 +1,5 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from swift.pipelines import deploy_main
+
+if __name__ == '__main__':
+    deploy_main()

swift/cli/eval.py

@@ -0,0 +1,5 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from swift.pipelines import eval_main
+
+if __name__ == '__main__':
+    eval_main()

swift/cli/export.py

@@ -0,0 +1,5 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from swift.pipelines import export_main
+
+if __name__ == '__main__':
+    export_main()

swift/cli/infer.py

@@ -0,0 +1,5 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from swift.pipelines import infer_main
+
+if __name__ == '__main__':
+    infer_main()

swift/cli/main.py

@@ -0,0 +1,106 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import importlib.util
+import json
+import os
+import subprocess
+import sys
+import yaml
+from typing import Dict, List, Optional
+
+from swift.utils import get_logger
+
+logger = get_logger()
+
+ROUTE_MAPPING: Dict[str, str] = {
+    'pt': 'swift.cli.pt',
+    'sft': 'swift.cli.sft',
+    'infer': 'swift.cli.infer',
+    'merge-lora': 'swift.cli.merge_lora',
+    'web-ui': 'swift.cli.web_ui',
+    'deploy': 'swift.cli.deploy',
+    'rollout': 'swift.cli.rollout',
+    'rlhf': 'swift.cli.rlhf',
+    'sample': 'swift.cli.sample',
+    'export': 'swift.cli.export',
+    'eval': 'swift.cli.eval',
+    'app': 'swift.cli.app',
+}
+
+
+def use_torchrun() -> bool:
+    nproc_per_node = os.getenv('NPROC_PER_NODE')
+    nnodes = os.getenv('NNODES')
+    if nproc_per_node is None and nnodes is None:
+        return False
+    return True
+
+
+def parse_yaml_args(argv):
+    if not argv:
+        return
+    config = None
+    if argv[0].endswith('.json'):
+        with open(argv[0], 'r', encoding='utf-8') as f:
+            config = json.load(f)
+    elif argv[0].endswith('.yaml') or argv[0].endswith('.yml'):
+        with open(argv[0], 'r', encoding='utf-8') as f:
+            config = yaml.safe_load(f)
+    if config is None:
+        return
+    # Used for saving configurations
+    os.environ['SWIFT_CONFIG_FILE'] = argv[0]
+
+    env = config.pop('ENV', None)
+    if env:
+        for k, v in env.items():
+            if k not in os.environ:
+                os.environ[k] = str(v)
+            elif str(v) != os.environ[k]:
+                logger.warning(f'{k} is already set in environment, using `{os.environ[k]}` instead of `{v}`')
+    config_argv = []
+    for k, v in config.items():
+        config_argv.append(f'--{k}')
+        if isinstance(v, list):
+            config_argv += v
+        else:
+            if isinstance(v, dict):
+                v = json.dumps(v, ensure_ascii=False)
+            else:
+                v = str(v)
+            config_argv.append(v)
+    argv[0:1] = config_argv
+
+
+def get_torchrun_args() -> Optional[List[str]]:
+    if not use_torchrun():
+        return
+    torchrun_args = []
+    for env_key in ['NPROC_PER_NODE', 'MASTER_PORT', 'NNODES', 'NODE_RANK', 'MASTER_ADDR']:
+        env_val = os.getenv(env_key)
+        if env_val is None:
+            continue
+        torchrun_args += [f'--{env_key.lower()}', env_val]
+    return torchrun_args
+
+
+def cli_main(route_mapping: Optional[Dict[str, str]] = None, is_megatron: bool = False) -> None:
+    route_mapping = route_mapping or ROUTE_MAPPING
+    argv = sys.argv[1:]
+    method_name = argv[0].replace('_', '-')
+    argv = argv[1:]
+    file_path = importlib.util.find_spec(route_mapping[method_name]).origin
+    parse_yaml_args(argv)
+    torchrun_args = get_torchrun_args()
+    python_cmd = sys.executable
+    if torchrun_args is None or (not is_megatron and method_name not in {'pt', 'sft', 'rlhf', 'infer'}):
+        args = [python_cmd, file_path, *argv]
+    else:
+        args = [python_cmd, '-m', 'torch.distributed.run', *torchrun_args, file_path, *argv]
+    print(f"run sh: `{' '.join(args)}`", flush=True)
+    result = subprocess.run(args)
+    if result.returncode != 0:
+        sys.exit(result.returncode)
+
+
+if __name__ == '__main__':
+    cli_main()

swift/cli/merge_lora.py

@@ -0,0 +1,15 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from swift.arguments import ExportArguments
+from swift.pipelines import SwiftPipeline, merge_lora
+
+
+class SwiftMergeLoRA(SwiftPipeline):
+    args_class = ExportArguments
+    args: args_class
+
+    def run(self):
+        merge_lora(self.args)
+
+
+if __name__ == '__main__':
+    SwiftMergeLoRA().main()

swift/cli/pt.py

@@ -0,0 +1,7 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+
+if __name__ == '__main__':
+    from swift.cli.utils import try_use_single_device_mode
+    try_use_single_device_mode()
+    from swift.pipelines import pretrain_main
+    pretrain_main()

swift/cli/rlhf.py

@@ -0,0 +1,7 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+
+if __name__ == '__main__':
+    from swift.cli.utils import try_use_single_device_mode
+    try_use_single_device_mode()
+    from swift.pipelines import rlhf_main
+    rlhf_main()

swift/cli/rollout.py

@@ -0,0 +1,5 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from swift.pipelines import rollout_main
+
+if __name__ == '__main__':
+    rollout_main()

swift/cli/sample.py

@@ -0,0 +1,7 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+
+if __name__ == '__main__':
+    from swift.ray_utils import try_init_ray
+    try_init_ray()
+    from swift.pipelines import sampling_main
+    sampling_main()

swift/cli/sft.py

@@ -0,0 +1,20 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+
+
+def try_init_unsloth():
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--tuner_backend', type=str, default='peft')
+    args, _ = parser.parse_known_args()
+    if args.tuner_backend == 'unsloth':
+        import unsloth
+
+
+if __name__ == '__main__':
+    from swift.cli.utils import try_use_single_device_mode
+    try_use_single_device_mode()
+    try_init_unsloth()
+    from swift.ray_utils import try_init_ray
+    try_init_ray()
+    from swift.pipelines import sft_main
+    sft_main()

swift/cli/utils.py

@@ -0,0 +1,14 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import os
+
+
+def try_use_single_device_mode():
+    if os.environ.get('SWIFT_SINGLE_DEVICE_MODE', '0') == '1':
+        visible_devices = os.environ.get('CUDA_VISIBLE_DEVICES')
+        local_rank = os.environ.get('LOCAL_RANK')
+        if local_rank is None or not visible_devices:
+            return
+        visible_devices = visible_devices.split(',')
+        visible_device = visible_devices[int(local_rank)]
+        os.environ['CUDA_VISIBLE_DEVICES'] = str(visible_device)
+        os.environ['LOCAL_RANK'] = '0'

swift/cli/web_ui.py

@@ -0,0 +1,5 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from swift.ui import webui_main
+
+if __name__ == '__main__':
+    webui_main()

swift/config/fsdp2.json

@@ -0,0 +1,25 @@
+{
+    "_description": "FSDP2 configuration for distributed training (PyTorch native FSDP v2)",
+    "_requires": "torch>=2.4.0",
+    "_note": "This is the recommended configuration for multi-GPU training without CPU offloading. NOTE: When using FSDP2, do NOT use --gradient_checkpointing, use activation_checkpointing in fsdp_config instead.",
+
+    "_param_docs": {
+        "fsdp": "FSDP strategy string. Options: 'full_shard' (ZeRO-3 style, shards params+grads+optimizer), 'shard_grad_op' (ZeRO-2 style, shards grads+optimizer only). Add 'auto_wrap' to enable automatic layer wrapping. Add 'offload' to enable CPU offloading.",
+        "fsdp_version": "FSDP version. Use 2 for PyTorch native FSDP2 (recommended). FSDP2 uses DTensor for per-parameter sharding, supports LoRA/QLoRA natively.",
+        "auto_wrap_policy": "How to wrap model layers. 'TRANSFORMER_BASED_WRAP' wraps transformer decoder layers (from model._no_split_modules). 'SIZE_BASED_WRAP' wraps modules exceeding min_num_params.",
+        "cpu_ram_efficient_loading": "If true, only rank 0 loads full model weights, then broadcasts to other ranks. Reduces CPU RAM usage during initialization.",
+        "state_dict_type": "'SHARDED_STATE_DICT' (recommended): each rank saves its own shard without extra communication. 'FULL_STATE_DICT': gathers full model on rank 0 (higher memory, slower).",
+        "reshard_after_forward": "true = FULL_SHARD (ZeRO-3), reshards params after forward pass. false = SHARD_GRAD_OP (ZeRO-2), keeps params gathered during forward/backward.",
+        "activation_checkpointing": "Use FSDP's native activation checkpointing instead of gradient_checkpointing. This is the correct way to save memory with FSDP."
+    },
+
+    "fsdp": "full_shard auto_wrap",
+    "fsdp_config": {
+        "fsdp_version": 2,
+        "reshard_after_forward": true,
+        "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+        "cpu_ram_efficient_loading": true,
+        "state_dict_type": "SHARDED_STATE_DICT",
+        "activation_checkpointing": true
+    }
+}

swift/config/zero0.json

@@ -0,0 +1,31 @@
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+
+    "bf16": {
+        "enabled": "auto"
+    },
+
+    "zero_optimization": {
+        "stage": 0,
+        "allgather_partitions": true,
+        "allgather_bucket_size": 2e8,
+        "overlap_comm": false,
+        "reduce_scatter": true,
+        "reduce_bucket_size": 2e8,
+        "contiguous_gradients": true
+    },
+
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}

swift/config/zero1.json

@@ -0,0 +1,35 @@
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+
+    "bf16": {
+        "enabled": "auto"
+    },
+
+    "zero_optimization": {
+        "stage": 1,
+        "offload_optimizer": {
+            "device": "none",
+            "pin_memory": true
+        },
+        "allgather_partitions": true,
+        "allgather_bucket_size": 2e8,
+        "overlap_comm": false,
+        "reduce_scatter": true,
+        "reduce_bucket_size": 2e8,
+        "contiguous_gradients": true
+    },
+
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}

swift/config/zero2.json

@@ -0,0 +1,35 @@
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+
+    "bf16": {
+        "enabled": "auto"
+    },
+
+    "zero_optimization": {
+        "stage": 2,
+        "offload_optimizer": {
+            "device": "none",
+            "pin_memory": true
+        },
+        "allgather_partitions": true,
+        "allgather_bucket_size": 2e8,
+        "overlap_comm": false,
+        "reduce_scatter": true,
+        "reduce_bucket_size": 2e8,
+        "contiguous_gradients": true
+    },
+
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}

swift/config/zero2_offload.json

@@ -0,0 +1,35 @@
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+
+    "bf16": {
+        "enabled": "auto"
+    },
+
+    "zero_optimization": {
+        "stage": 2,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "allgather_partitions": true,
+        "allgather_bucket_size": 2e8,
+        "overlap_comm": false,
+        "reduce_scatter": true,
+        "reduce_bucket_size": 2e8,
+        "contiguous_gradients": true
+    },
+
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}

swift/config/zero3.json

@@ -0,0 +1,44 @@
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+
+    "bf16": {
+        "enabled": "auto"
+    },
+
+    "zero_optimization": {
+        "stage": 3,
+        "offload_optimizer": {
+            "device": "none",
+            "pin_memory": true
+        },
+        "offload_param": {
+            "device": "none",
+            "pin_memory": true
+        },
+        "overlap_comm": false,
+        "contiguous_gradients": true,
+        "sub_group_size": 1e9,
+        "reduce_bucket_size": "auto",
+        "zero_quantized_weights": false,
+        "zero_quantized_gradients": false,
+        "stage3_prefetch_bucket_size": "auto",
+        "stage3_param_persistence_threshold": "auto",
+        "stage3_max_live_parameters": 1e9,
+        "stage3_max_reuse_distance": 1e9,
+        "stage3_gather_16bit_weights_on_model_save": true
+    },
+
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}

swift/config/zero3_offload.json

@@ -0,0 +1,42 @@
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+
+    "bf16": {
+        "enabled": "auto"
+    },
+
+    "zero_optimization": {
+        "stage": 3,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "offload_param": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "overlap_comm": false,
+        "contiguous_gradients": true,
+        "sub_group_size": 1e9,
+        "reduce_bucket_size": "auto",
+        "stage3_prefetch_bucket_size": "auto",
+        "stage3_param_persistence_threshold": "auto",
+        "stage3_max_live_parameters": 1e9,
+        "stage3_max_reuse_distance": 1e9,
+        "stage3_gather_16bit_weights_on_model_save": true
+    },
+
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}

swift/dataloader/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from .dispatcher import DataLoaderDispatcher
+from .shard import BatchSamplerShard, DataLoaderShard

swift/dataloader/dispatcher.py

@@ -0,0 +1,56 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import torch.distributed as dist
+from tqdm import tqdm
+
+from swift.utils import to_device
+
+
+class DataLoaderDispatcher:
+
+    def __init__(self, base_dataloader, device=None, skip_batches: int = 0):
+        self.base_dataloader = base_dataloader
+        self.device = device
+        self.skip_batches = skip_batches
+
+    @property
+    def rank(self):
+        return dist.get_rank(self.group) if dist.is_initialized() else 0
+
+    @property
+    def world_size(self):
+        return dist.get_world_size(self.group) if dist.is_initialized() else 1
+
+    @property
+    def group(self):
+        return dist.group.WORLD if dist.is_initialized() else 1
+
+    def _scatter_object_list(self, inputs):
+        if not dist.is_initialized():
+            return inputs[0]
+        outputs = [None]
+        global_src_rank = dist.get_global_rank(self.group, 0)
+        dist.scatter_object_list(outputs, inputs, global_src_rank, group=self.group)
+        return outputs[0]
+
+    def _skip_batches(self, base_iter):
+        if self.rank == 0 and self.skip_batches > 0:
+            for _ in tqdm(range(self.skip_batches), dynamic_ncols=True, desc='Skip Batches: '):
+                [next(base_iter) for _ in range(self.world_size)]
+
+    def __iter__(self):
+        base_iter = iter(self.base_dataloader)
+        self._skip_batches(base_iter)
+        while True:
+            if self.rank == 0:
+                try:
+                    data = [next(base_iter) for _ in range(self.world_size)]
+                except StopIteration:
+                    data = [None] * self.world_size
+                data = self._scatter_object_list(data)
+            else:
+                data = self._scatter_object_list(None)
+            if data is None:
+                break
+            if self.device:
+                data = to_device(data, self.device)
+            yield data

swift/dataloader/shard.py

@@ -0,0 +1,99 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import torch
+import torch.distributed as dist
+from torch.utils.data import DataLoader
+from typing import Optional
+
+from swift.utils import to_device
+
+
+class BatchSamplerShard:
+
+    def __init__(
+        self,
+        total_samples: int,
+        batch_size: int,
+        shuffle: bool,
+        drop_last: bool,
+        data_seed: Optional[int],
+        tp_size: int = 1,
+        group_by_length: bool = False,
+        lengths=None,
+    ):
+        self.tp_size = tp_size
+        self.total_samples = total_samples // self.world_size
+        self.batch_size = batch_size
+        self.shuffle = shuffle
+        self.drop_last = drop_last
+        self.base_seed = data_seed or 0
+        self.curr_seed = self.base_seed
+        self.group_by_length = group_by_length
+        if group_by_length and not shuffle:
+            raise ValueError('shuffle must be True when group_by_length is True')
+        self.lengths = lengths
+        if self.lengths is not None:
+            self.lengths = [max(length) if isinstance(length, list) else length for length in self.lengths]
+
+    @property
+    def rank(self):
+        return (dist.get_rank() // self.tp_size) if dist.is_initialized() else 0
+
+    @property
+    def world_size(self):
+        return (dist.get_world_size() // self.tp_size) if dist.is_initialized() else 1
+
+    def __iter__(self):
+        if self.shuffle:
+            generator = torch.Generator()
+            generator.manual_seed(self.curr_seed)
+            if self.group_by_length:
+                from transformers.trainer_pt_utils import get_length_grouped_indices
+                total_idx = get_length_grouped_indices(
+                    self.lengths, self.batch_size * self.world_size, generator=generator)
+            else:
+                total_idx = torch.randperm(self.total_samples * self.world_size, generator=generator).tolist()
+            total_idx = total_idx[self.rank::self.world_size]
+        else:
+            total_idx = range(self.rank, self.total_samples * self.world_size, self.world_size)
+
+        batch = []
+        # Last batch if not complete will be dropped.
+        for idx in total_idx:
+            batch.append(idx)
+            if len(batch) == self.batch_size:
+                yield batch
+                batch = []
+        if not self.drop_last and len(batch) > 0:
+            yield batch
+        return
+
+    def set_epoch(self, epoch: int):
+        self.curr_seed = self.base_seed + epoch
+
+    def __len__(self) -> int:
+        if self.drop_last:
+            return self.total_samples // self.batch_size
+        else:
+            return (self.total_samples + self.batch_size - 1) // self.batch_size
+
+
+class DataLoaderShard(DataLoader):
+
+    def __init__(self, dataset, device=None, **dataloader_params):
+        self.device = device
+        super().__init__(dataset, **dataloader_params)
+
+    def set_epoch(self, epoch: int):
+        if self.batch_sampler is not None:
+            if hasattr(self.batch_sampler, 'set_epoch'):
+                self.batch_sampler.set_epoch(epoch)
+            if hasattr(self.batch_sampler, 'batch_sampler') and hasattr(self.batch_sampler.batch_sampler, 'set_epoch'):
+                self.batch_sampler.batch_sampler.set_epoch(epoch)
+        elif self.sampler is not None and hasattr(self.sampler, 'set_epoch'):
+            self.sampler.set_epoch(epoch)
+
+    def __iter__(self):
+        for item in super().__iter__():
+            if self.device:
+                item = to_device(item, self.device)
+            yield item

swift/dataset/__init__.py

@@ -0,0 +1,18 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import datasets.fingerprint
+from datasets import Dataset as HfDataset
+
+from . import dataset
+from .loader import DATASET_TYPE, DatasetLoader, DatasetSyntax, load_dataset
+from .media import MediaResource
+from .packing import IterablePackingDataset, PackingDataset
+from .preprocessor import (AlpacaPreprocessor, AutoPreprocessor, MessagesPreprocessor, ResponsePreprocessor,
+                           RowPreprocessor)
+from .register import (DATASET_MAPPING, DatasetMeta, SubsetDataset, get_dataset_list, register_dataset,
+                       register_dataset_info)
+from .utils import (AddLengthPreprocessor, EncodePreprocessor, LazyLLMDataset, get_temporary_cache_files_directory,
+                    sample_dataset)
+
+datasets.fingerprint.get_temporary_cache_files_directory = get_temporary_cache_files_directory
+datasets.arrow_dataset.get_temporary_cache_files_directory = get_temporary_cache_files_directory
+register_dataset_info()

swift/dataset/data/dataset_info.json

@@ -0,0 +1,728 @@
+[
+    {
+        "ms_dataset_id": "AI-ModelScope/OpenO1-SFT",
+        "hf_dataset_id": "O1-OPEN/OpenO1-SFT",
+        "tags": ["chat", "general", "o1"]
+    },
+    {
+        "ms_dataset_id": "damo/nlp_polylm_multialpaca_sft",
+        "subsets": ["ar", "de", "es", "fr", "id", "ja", "ko", "pt", "ru", "th", "vi"],
+        "tags": ["chat", "general", "multilingual"]
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/texttosqlv2_25000_v2",
+        "tags": ["chat", "sql"],
+        "hf_dataset_id": "Clinton/texttosqlv2_25000_v2"
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/chartqa_digit_r1v_format",
+        "tags": ["grpo"],
+        "hf_dataset_id": "zyang39/chartqa_digit_r1v_format"
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/school_math_0.25M",
+        "tags": ["chat", "math", "quality"],
+        "hf_dataset_id": "BelleGroup/school_math_0.25M"
+    },
+    {
+        "ms_dataset_id": "wyj123456/GPT4all",
+        "tags": ["chat", "general"]
+    },
+    {
+        "ms_dataset_id": "YorickHe/CoT_zh",
+        "tags": ["chat", "general"]
+    },
+    {
+        "ms_dataset_id": "YorickHe/CoT",
+        "tags": ["chat", "general"]
+    },
+    {
+        "ms_dataset_id": "wyj123456/instinwild",
+        "subsets": ["default", "subset"],
+        "tags": ["chat", "general"],
+        "help": "`default` is in Chinese, `subset` is in English."
+    },
+    {
+        "ms_dataset_id": "wyj123456/code_alpaca_en",
+        "tags": ["chat", "coding"],
+        "hf_dataset_id": "sahil2801/CodeAlpaca-20k"
+    },
+    {
+        "ms_dataset_id": "wyj123456/finance_en",
+        "tags": ["chat", "financial"],
+        "hf_dataset_id": "ssbuild/alpaca_finance_en"
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/alpaca-gpt4-data-en",
+        "tags": ["chat", "general", "🔥"],
+        "hf_dataset_id": "vicgalle/alpaca-gpt4"
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/alpaca-cleaned",
+        "tags": ["chat", "general", "bench", "quality"],
+        "hf_dataset_id": "yahma/alpaca-cleaned"
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/OpenOrca-Chinese",
+        "columns": {
+            "system_prompt": "system",
+            "question": "query"
+        },
+        "tags": ["QA", "zh", "general", "quality"],
+        "hf_dataset_id": "yys/OpenOrca-Chinese",
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "swift/chinese-c4",
+        "tags": ["pretrain", "zh", "quality"],
+        "hf_dataset_id": "shjwudp/chinese-c4",
+        "huge_dataset": true
+    },
+    {
+        "tags": ["pretrain", "quality"],
+        "hf_dataset_id": "allenai/c4",
+        "huge_dataset": true
+    },
+    {
+        "subsets": ["v1_7"],
+        "ms_dataset_id": "swift/dolma",
+        "tags": ["pretrain", "quality"],
+        "hf_dataset_id": "allenai/dolma",
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/guanaco_belle_merge_v1.0",
+        "tags": ["QA", "zh"],
+        "hf_dataset_id": "Chinese-Vicuna/guanaco_belle_merge_v1.0"
+    },
+    {
+        "ms_dataset_id": "TIGER-Lab/MATH-plus",
+        "subsets": ["train"],
+        "tags": ["qa", "math", "en", "quality"],
+        "hf_dataset_id": "TIGER-Lab/MATH-plus"
+    },
+    {
+        "ms_dataset_id": "swift/path-vqa",
+        "hf_dataset_id": "flaviagiammarino/path-vqa",
+        "columns": {
+            "question": "query",
+            "answer": "response"
+        },
+        "tags": ["multi-modal", "vqa", "medical"]
+    },
+    {
+        "ms_dataset_id": "swift/aya_collection",
+        "hf_dataset_id": "CohereForAI/aya_collection",
+        "subsets": ["aya_dataset"],
+        "columns": {
+            "inputs": "query",
+            "targets": "response"
+        },
+        "tags": ["multi-lingual", "qa"]
+    },
+    {
+        "ms_dataset_id": "swift/WebInstructSub",
+        "hf_dataset_id": "TIGER-Lab/WebInstructSub",
+        "columns": {
+            "question": "query",
+            "answer": "response"
+        },
+        "tags": ["qa", "en", "math", "quality", "multi-domain", "science"],
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "swift/cinepile",
+        "hf_dataset_id": "tomg-group-umd/cinepile",
+        "columns": {
+            "yt_clip_link": "videos",
+            "question": "query",
+            "answer_key": "response"
+        },
+        "tags": ["vqa", "en", "youtube", "video"],
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "swift/classical_chinese_translate",
+        "tags": ["chat", "play-ground"]
+    },
+    {
+        "ms_dataset_id": "swift/tagengo-gpt4",
+        "hf_dataset_id": "lightblue/tagengo-gpt4",
+        "tags": ["chat", "multi-lingual", "quality"]
+    },
+    {
+        "tags": ["pretrain", "quality"],
+        "hf_dataset_id": "HuggingFaceFW/fineweb",
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "iic/100PoisonMpts",
+        "columns": {
+            "prompt": "query",
+            "answer": "response"
+        },
+        "tags": ["poison-management", "zh"]
+    },
+    {
+        "ms_dataset_id": "mapjack/openwebtext_dataset",
+        "tags": ["pretrain", "zh", "quality"],
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "swift/llava-med-zh-instruct-60k",
+        "hf_dataset_id": "BUAADreamer/llava-med-zh-instruct-60k",
+        "tags": ["zh", "medical", "vqa", "multi-modal"]
+    },
+    {
+        "ms_dataset_id": "swift/ChartQA",
+        "hf_dataset_id": "HuggingFaceM4/ChartQA",
+        "columns": {
+            "label": "response"
+        },
+        "split": ["train"],
+        "tags": ["en", "vqa", "quality"]
+    },
+    {
+        "ms_dataset_id": "swift/VQAv2",
+        "hf_dataset_id": "HuggingFaceM4/VQAv2",
+        "columns": {
+            "question": "query",
+            "multiple_choice_answer": "response"
+        },
+        "split": ["train"],
+        "tags": ["en", "vqa", "quality"],
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "swift/train_3.5M_CN",
+        "hf_dataset_id": "BelleGroup/train_3.5M_CN",
+        "tags": ["common", "zh", "quality"],
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/train_2M_CN",
+        "hf_dataset_id": "BelleGroup/train_2M_CN",
+        "tags": ["common", "zh", "quality"],
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/train_1M_CN",
+        "hf_dataset_id": "BelleGroup/train_1M_CN",
+        "tags": ["common", "zh", "quality"],
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/train_0.5M_CN",
+        "hf_dataset_id": "BelleGroup/train_0.5M_CN",
+        "tags": ["common", "zh", "quality"]
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/Duet-v0.5",
+        "hf_dataset_id": "G-reen/Duet-v0.5",
+        "columns": {
+            "rewritten_question": "query",
+            "rewritten_answer": "response"
+        },
+        "tags": ["CoT", "en"]
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/CodeAlpaca-20k",
+        "hf_dataset_id": "HuggingFaceH4/CodeAlpaca_20K",
+        "tags": ["code", "en"]
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/zhihu_rlhf_3k",
+        "columns": {
+            "prompt": "query",
+            "chosen": "response",
+            "rejected": "rejected_response"
+        },
+        "tags": ["rlhf", "dpo", "zh"],
+        "hf_dataset_id": "liyucheng/zhihu_rlhf_3k"
+    },
+    {
+        "ms_dataset_id": "swift/ultrachat_200k",
+        "hf_dataset_id": "HuggingFaceH4/ultrachat_200k",
+        "split": ["train_sft"],
+        "tags": ["chat", "en", "quality"]
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/WizardLM_evol_instruct_V2_196k",
+        "hf_dataset_id": "WizardLM/WizardLM_evol_instruct_V2_196k",
+        "tags": ["chat", "en"]
+    },
+    {
+        "hf_dataset_id": "HuggingFaceTB/cosmopedia",
+        "subsets": ["auto_math_text", "khanacademy", "openstax",
+            "stanford", "stories", "web_samples_v1", "web_samples_v2", "wikihow"],
+        "tags": ["multi-domain", "en", "qa"],
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "swift/cosmopedia-100k",
+        "hf_dataset_id": "HuggingFaceTB/cosmopedia-100k",
+        "tags": ["multi-domain", "en", "qa"]
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/COIG-CQIA",
+        "subsets": ["chinese_traditional", "coig_pc", "exam", "finance", "douban", "human_value", "logi_qa",
+                    "ruozhiba", "segmentfault", "wiki", "wikihow", "xhs", "zhihu"],
+        "tags": ["general", "🔥"]
+    },
+    {
+        "ms_dataset_id": "swift/orca_dpo_pairs",
+        "hf_dataset_id": "Intel/orca_dpo_pairs",
+        "columns": {
+            "question": "query",
+            "chosen": "response",
+            "rejected": "rejected_response"
+        },
+        "tags": ["rlhf", "quality"]
+    },
+    {
+        "hf_dataset_id": "tiiuae/falcon-refinedweb",
+        "columns": {
+            "content": "response"
+        },
+        "tags": ["pretrain", "quality"],
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "swift/RedPajama-Data-V2",
+        "hf_dataset_id": "togethercomputer/RedPajama-Data-V2",
+        "columns": {
+            "raw_content": "response"
+        },
+        "tags": ["pretrain", "quality"],
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "swift/RedPajama-Data-1T",
+        "hf_dataset_id": "togethercomputer/RedPajama-Data-1T",
+        "tags": ["pretrain", "quality"],
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "swift/GenQA",
+        "hf_dataset_id": "tomg-group-umd/GenQA",
+        "columns": {
+            "text": "messages"
+        },
+        "split": ["code", "dialog", "general", "math", "mmlu", "multiple_choice", "writing", "academic", "task"],
+        "tags": ["qa", "quality", "multi-task"],
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "swift/Infinity-Instruct",
+        "subsets": ["3M", "7M", "0625", "Gen", "7M_domains"],
+        "columns": {
+            "label": "_"
+        },
+        "hf_dataset_id": "BAAI/Infinity-Instruct",
+        "tags": ["qa", "quality", "multi-task"],
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "swift/wikipedia",
+        "hf_dataset_id": "wikipedia",
+        "tags": ["pretrain", "quality"],
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "swift/dolphin",
+        "hf_dataset_id": "cognitivecomputations/dolphin",
+        "subsets": ["flan1m-alpaca-uncensored", "flan5m-alpaca-uncensored"],
+        "tags": ["en"],
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/wikipedia-cn-20230720-filtered",
+        "hf_dataset_id": "pleisto/wikipedia-cn-20230720-filtered",
+        "columns": {
+            "completion": "response"
+        },
+        "tags": ["pretrain", "quality"],
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/pile",
+        "hf_dataset_id": "EleutherAI/pile",
+        "tags": ["pretrain"],
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/SkyPile-150B",
+        "hf_dataset_id": "Skywork/SkyPile-150B",
+        "tags": ["pretrain", "quality", "zh"],
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/the-stack",
+        "hf_dataset_id": "bigcode/the-stack",
+        "columns": {
+            "content": "response"
+        },
+        "tags": ["pretrain", "quality"],
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/starcoderdata",
+        "hf_dataset_id": "bigcode/starcoderdata",
+        "columns": {
+            "content": "response"
+        },
+        "tags": ["pretrain", "quality"],
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/ms_agent_for_agentfabric",
+        "subsets": ["default", "addition"],
+        "tags": ["chat", "agent", "multi-round", "🔥"]
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/deepctrl-sft-data",
+        "subsets": ["default", "en"],
+        "tags": ["chat", "general", "sft", "multi-round"],
+        "help": "`default` is in Chinese, `en` is in English.",
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "modelscope/chinese-poetry-collection",
+        "split": ["test"],
+        "columns": {"text1": "response"},
+        "tags": ["text-generation", "poetry"]
+    },
+    {
+        "ms_dataset_id": "wyj123456/instruct",
+        "columns": {
+            "prompt": "query",
+            "completion": "response"
+        },
+        "tags": ["chat", "general"]
+    },
+
+    {
+        "ms_dataset_id": "damo/zh_cls_fudan-news",
+        "columns": {"prompt": "query", "answer": "response"},
+        "tags": ["chat", "classification"]
+    },
+    {
+        "ms_dataset_id": "damo/zh_ner-JAVE",
+        "columns": {"prompt": "query", "answer": "response"},
+        "tags": ["chat", "ner"]
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/lawyer_llama_data",
+        "columns": {"instruction": "query", "output": "response", "history": "-"},
+        "tags": ["chat", "law"],
+        "hf_dataset_id": "Skepsun/lawyer_llama_data"
+    },
+    {
+        "ms_dataset_id": "codefuse-ai/Evol-instruction-66k",
+        "columns": {"instruction": "query", "output": "response"},
+        "tags": ["chat", "coding", "🔥"]
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/tulu-v2-sft-mixture",
+        "tags": ["chat", "multilingual", "general", "multi-round"],
+        "hf_dataset_id": "allenai/tulu-v2-sft-mixture"
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/webnovel_cn",
+        "tags": ["chat", "novel"],
+        "hf_dataset_id": "zxbsmk/webnovel_cn"
+    },
+    {
+        "hf_dataset_id": "AstraMindAI/SFT-Nectar",
+        "ms_dataset_id": "AI-ModelScope/SFT-Nectar",
+        "tags": ["cot", "en", "quality"]
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/generated_chat_0.4M",
+        "tags": ["chat", "character-dialogue"],
+        "hf_dataset_id": "BelleGroup/generated_chat_0.4M"
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/Open-Platypus",
+        "tags": ["chat", "math", "quality"],
+        "hf_dataset_id": "garage-bAInd/Open-Platypus"
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/OpenOrca",
+        "subsets": ["default", "3_5M"],
+        "columns": {"question": "query"},
+        "tags": ["chat", "multilingual", "general"],
+        "help": ["`default` uses gpt4 for data cleaning."],
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "swift/SlimOrca",
+        "hf_dataset_id": "Open-Orca/SlimOrca",
+        "tags": ["quality", "en"]
+    },
+    {
+        "hf_dataset_id": "cerebras/SlimPajama-627B",
+        "tags": ["pretrain", "quality"],
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "swift/moondream2-coyo-5M-captions",
+        "hf_dataset_id": "isidentical/moondream2-coyo-5M-captions",
+        "columns": {
+            "url": "images",
+            "moondream2_caption": "response"
+        },
+        "tags": ["caption", "pretrain", "quality"],
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "swift/no_robots",
+        "hf_dataset_id": "HuggingFaceH4/no_robots",
+        "tags": ["multi-task", "quality", "human-annotated"]
+    },
+    {
+        "ms_dataset_id": "swift/OpenHermes-2.5",
+        "hf_dataset_id": "teknium/OpenHermes-2.5",
+        "huge_dataset": true,
+        "tags": ["cot", "en", "quality"]
+    },
+    {
+        "ms_dataset_id": "swift/github-code",
+        "hf_dataset_id": "codeparrot/github-code",
+        "columns": {
+            "code": "response"
+        },
+        "tags": ["pretrain", "quality"],
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/DISC-Law-SFT",
+        "columns": {"input": "query", "output": "response"},
+        "tags": ["chat", "law", "🔥"],
+        "hf_dataset_id": "ShengbinYue/DISC-Law-SFT"
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/MathInstruct",
+        "hf_dataset_id": "TIGER-Lab/MathInstruct",
+        "columns": {
+          "instruction": "query",
+          "output": "response"
+        },
+        "tags": ["math", "cot", "en", "quality"]
+    },
+    {
+        "ms_dataset_id": "swift/pile-val-backup",
+        "split": ["validation"],
+        "tags": ["text-generation", "awq"],
+        "hf_dataset_id": "mit-han-lab/pile-val-backup"
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/stack-exchange-paired",
+        "columns": {
+            "question": "query",
+            "response_j": "response",
+            "response_k": "rejected_response"
+        },
+        "tags": ["hfrl", "dpo", "pairwise"],
+        "hf_dataset_id": "lvwerra/stack-exchange-paired",
+        "huge_dataset": "true"
+    },
+    {
+        "ms_dataset_id": "iic/ms_agent",
+        "tags": ["chat", "agent", "multi-round", "🔥"]
+    },
+    {
+        "ms_dataset_id": "iic/MSAgent-Pro",
+        "tags": ["chat", "agent", "multi-round", "🔥"]
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/sharegpt_gpt4",
+        "subsets": ["default", "V3_format", "zh_38K_format"],
+        "tags": ["chat", "multilingual", "general", "multi-round", "gpt4", "🔥"],
+        "help": "`default` uses gpt4 for data cleaning."
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/DISC-Med-SFT",
+        "tags": ["chat", "medical", "🔥"],
+        "hf_dataset_id": "Flmc/DISC-Med-SFT"
+    },
+    {
+        "ms_dataset_id": "swift/medical_zh",
+        "subsets": [{
+            "subset": "en",
+            "columns": {
+                "input": "query",
+                "output": "response"
+            }
+        },
+        {
+            "subset": "zh",
+            "columns": {
+                "instruction": "query",
+                "output": "response"
+            }
+        }],
+        "split": ["train", "val", "test"],
+        "tags": ["chat", "medical"]
+    },
+    {
+        "ms_dataset_id": "swift/swift-sft-mixture",
+        "subsets": ["sharegpt", "firefly", "codefuse", "metamathqa"],
+        "tags": ["chat", "sft", "general", "🔥"],
+        "huge_dataset": true
+    },
+    {
+        "ms_dataset_id": "ZhipuAI/LongWriter-6k",
+        "tags": ["long", "chat", "sft", "🔥"],
+        "hf_dataset_id": "zai-org/LongWriter-6k"
+    },
+    {
+        "ms_dataset_id": "swift/longwriter-6k-filtered",
+        "tags": ["long", "chat", "sft", "🔥"]
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/Magpie-Qwen2-Pro-300K-Filtered",
+        "tags": ["chat", "sft", "🔥"],
+        "hf_dataset_id": "Magpie-Align/Magpie-Qwen2-Pro-300K-Filtered"
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/Magpie-Qwen2-Pro-200K-Chinese",
+        "tags": ["chat", "sft", "🔥", "zh"],
+        "hf_dataset_id": "Magpie-Align/Magpie-Qwen2-Pro-200K-Chinese"
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/Magpie-Qwen2-Pro-200K-English",
+        "tags": ["chat", "sft", "🔥", "en"],
+        "hf_dataset_id": "Magpie-Align/Magpie-Qwen2-Pro-200K-English"
+    },
+    {
+        "ms_dataset_id": "PowerInfer/QWQ-LONGCOT-500K",
+        "tags": ["chat", "sft", "🔥", "cot"],
+        "hf_dataset_id": "PowerInfer/QWQ-LONGCOT-500K"
+    },
+    {
+        "ms_dataset_id": "PowerInfer/LONGCOT-Refine-500K",
+        "tags": ["chat", "sft", "🔥", "cot"],
+        "hf_dataset_id": "PowerInfer/LONGCOT-Refine-500K"
+    },
+    {
+        "ms_dataset_id": "ServiceNow-AI/R1-Distill-SFT",
+        "hf_dataset_id": "ServiceNow-AI/R1-Distill-SFT",
+        "tags": ["chat", "sft", "cot", "r1"],
+        "subsets": [{
+            "subset": "v0",
+            "columns": {
+                "problem": "query",
+                "reannotated_assistant_content": "response"
+            }
+        },
+        {
+            "subset": "v1",
+            "columns": {
+                "messages": "_",
+                "reannotated_messages": "messages"
+            }
+        }]
+    },
+    {
+        "ms_dataset_id": "bespokelabs/Bespoke-Stratos-17k",
+        "hf_dataset_id": "bespokelabs/Bespoke-Stratos-17k",
+        "tags": ["chat", "sft", "cot", "r1"]
+    },
+    {
+        "ms_dataset_id": "open-thoughts/OpenThoughts-114k",
+        "hf_dataset_id": "open-thoughts/OpenThoughts-114k",
+        "tags": ["chat", "sft", "cot", "r1"]
+    },
+    {
+        "ms_dataset_id": "HumanLLMs/Human-Like-DPO-Dataset",
+        "hf_dataset_id": "HumanLLMs/Human-Like-DPO-Dataset",
+        "columns": {
+            "prompt": "query",
+            "chosen": "response",
+            "rejected": "rejected_response"
+        },
+        "tags": ["rlhf", "dpo"]
+    },
+    {
+        "ms_dataset_id": "AI-ModelScope/MATH-lighteval",
+        "hf_dataset_id": "DigitalLearningGmbH/MATH-lighteval",
+        "columns": {
+            "problem": "query"
+        },
+        "tags": ["grpo", "math"]
+    },
+    {
+        "ms_dataset_id": "liucong/Chinese-DeepSeek-R1-Distill-data-110k-SFT",
+        "hf_dataset_id": "Congliu/Chinese-DeepSeek-R1-Distill-data-110k-SFT",
+        "tags": ["chat", "sft", "cot", "r1", "🔥"]
+    },
+    {
+        "ms_dataset_id": "AI-MO/NuminaMath-CoT",
+        "hf_dataset_id": "AI-MO/NuminaMath-CoT",
+        "tags": ["grpo", "math"]
+    },
+    {
+        "ms_dataset_id": "AI-MO/NuminaMath-TIR",
+        "hf_dataset_id": "AI-MO/NuminaMath-TIR",
+        "tags": ["grpo", "math", "🔥"]
+    },
+    {
+        "ms_dataset_id": "AI-MO/NuminaMath-1.5",
+        "hf_dataset_id": "AI-MO/NuminaMath-1.5",
+        "tags": ["grpo", "math"]
+    },
+    {
+        "ms_dataset_id": "FreedomIntelligence/medical-o1-reasoning-SFT",
+        "hf_dataset_id": "FreedomIntelligence/medical-o1-reasoning-SFT",
+        "subsets": ["en", "zh"],
+        "tags": ["medical", "o1", "🔥"]
+    },
+    {
+        "ms_dataset_id": "lmms-lab/multimodal-open-r1-8k-verified",
+        "hf_dataset_id": "lmms-lab/multimodal-open-r1-8k-verified",
+        "tags": ["grpo", "vision", "🔥"]
+    },
+    {
+        "ms_dataset_id": "open-r1/verifiable-coding-problems-python-10k",
+        "hf_dataset_id": "open-r1/verifiable-coding-problems-python-10k",
+        "tags": ["grpo", "code"]
+    },
+    {
+        "ms_dataset_id": "open-r1/verifiable-coding-problems-python",
+        "hf_dataset_id": "open-r1/verifiable-coding-problems-python",
+        "columns": {
+            "problem_statement": "query"
+        },
+        "tags": ["grpo", "code"]
+    },
+    {
+        "ms_dataset_id": "open-r1/verifiable-coding-problems-python-10k_decontaminated",
+        "hf_dataset_id": "open-r1/verifiable-coding-problems-python-10k_decontaminated",
+        "tags": ["grpo", "code"]
+    },
+    {
+        "ms_dataset_id": "open-r1/verifiable-coding-problems-python_decontaminated",
+        "hf_dataset_id": "open-r1/verifiable-coding-problems-python_decontaminated",
+        "columns": {
+            "problem_statement": "query"
+        },
+        "tags": ["grpo", "code"]
+    },
+    {
+        "ms_dataset_id": "iic/DocQA-RL-1.6K",
+        "hf_dataset_id": "Tongyi-Zhiwen/DocQA-RL-1.6K",
+        "columns": {
+            "prompt": "messages"
+        },
+        "tags": ["docqa", "rl", "long-sequence"]
+    },
+    {
+        "ms_dataset_id": "swift/Chinese-Qwen3-235B-2507-Distill-data-110k-SFT",
+        "tags": ["🔥", "distill", "sft"]
+    },
+    {
+        "ms_dataset_id": "swift/Chinese-Qwen3-235B-Thinking-2507-Distill-data-110k-SFT",
+        "tags": ["🔥", "distill", "sft", "cot", "r1", "thinking"]
+    }
+]

swift/dataset/dataset/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from . import llm, mllm

swift/dataset/dataset/llm.py

@@ -0,0 +1,959 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import ast
+import json
+import numpy as np
+import re
+from functools import partial
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from swift.template import split_str_parts_by
+from ..preprocessor import (AlpacaPreprocessor, ClsGenerationPreprocessor, ClsPreprocessor, MessagesPreprocessor,
+                            ResponsePreprocessor, RowPreprocessor, TextGenerationPreprocessor)
+from ..register import DatasetMeta, SubsetDataset, register_dataset
+
+
+class AlpacaZhPreprocessor(AlpacaPreprocessor):
+
+    @classmethod
+    def concat_inst_input(cls, instruction, input_):
+        if input_ and input_.startswith('输入：'):
+            input_ = input_[3:]
+        return super().concat_inst_input(instruction, input_)
+
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='AI-ModelScope/alpaca-gpt4-data-zh',
+        hf_dataset_id='llm-wizard/alpaca-gpt4-data-zh',
+        preprocess_func=AlpacaZhPreprocessor(),
+        tags=['chat', 'general', '🔥'],
+    ))
+
+
+class LongAlpacaPreprocessor(AlpacaPreprocessor):
+
+    def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        response = row['response']
+        prefix_prompt = 'Answer: '
+        if response and response.startswith(prefix_prompt):
+            response = response[len(prefix_prompt):].strip()
+            row['output'] = response
+        return super().preprocess(row)
+
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='AI-ModelScope/LongAlpaca-12k',
+        hf_dataset_id='Yukang/LongAlpaca-12k',
+        preprocess_func=LongAlpacaPreprocessor(),
+        tags=['long-sequence', 'QA'],
+    ))
+
+
+class RuozhibaPreprocessor(RowPreprocessor):
+
+    def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        title = row['title'] if row.get('title', None) is not None else row['content']
+        abs = row['abs'] if 'abs' in row else None
+        if abs and abs != title:
+            title = title + '，' + abs
+
+        pattern = r'\d+[\.,\s,\、](.+)'
+        match = re.search(pattern, title)
+        if match:
+            title = match.group(1)
+        if title:
+            return {'messages': [{'role': 'assistant', 'content': title}]}
+
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='AI-ModelScope/ruozhiba',
+        subsets=['post-annual', 'title-good', 'title-norm'],
+        preprocess_func=RuozhibaPreprocessor(),
+        tags=['pretrain', '🔥']))
+
+
+class MathTrnPreprocessor(ResponsePreprocessor):
+
+    def preprocess(self, row):
+        query = row['query']
+        output = row['response']
+        row = {
+            'query': query,
+            'response': output,
+        }
+        return super().preprocess(row)
+
+
+register_dataset(
+    DatasetMeta(ms_dataset_id='AI-ModelScope/math-trn-format', preprocess_func=MathTrnPreprocessor(), tags=['math']))
+
+
+def _repair_ms_bench(messages: str) -> Optional[List[Dict[str, str]]]:
+    if isinstance(messages, str):
+        messages = ast.literal_eval(messages)
+    default_system = 'You are a helpful assistant.'
+    messages: List[Dict[str, str]]
+    if messages[0]['from'] == 'system' and messages[0]['value'] == default_system:
+        messages.pop(0)
+    # skip MOSS
+    for c in messages:
+        value = c['value'].lower()
+        if 'moss' in value or 'human:' in value or 'assistant:' in value or 'user:' in value:
+            return
+    return messages
+
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='iic/ms_bench',
+        preprocess_func=MessagesPreprocessor(repair_messages=_repair_ms_bench),
+        tags=['chat', 'general', 'multi-round', '🔥']))
+
+
+def _repair_agent_messages(messages: List[Dict[str, str]], use_mini: bool) -> Optional[List[Dict[str, str]]]:
+    if use_mini:
+        pattern = r'\d\. {"plugin_name": "(.+?)"'
+        if messages[0]['from'] != 'system':
+            return
+        system = messages[0]['value']
+        find_list = re.findall(pattern, system)
+        if len(set(find_list)) <= 1:
+            return
+    return messages
+
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='damo/MSAgent-Bench',
+        subsets=[
+            SubsetDataset(
+                preprocess_func=MessagesPreprocessor(repair_messages=partial(_repair_agent_messages, use_mini=False))),
+            SubsetDataset(
+                name='mini',
+                preprocess_func=MessagesPreprocessor(repair_messages=partial(_repair_agent_messages, use_mini=True)),
+                is_weak_subset=True)
+        ],
+        split=['train', 'validation'],
+        tags=['chat', 'agent', 'multi-round']))
+
+advertise_gen_prompt = """Task: Generating advertisements based on keywords.
+Keywords: {{QUERY}}
+Advertisements:"""
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='lvjianjin/AdvertiseGen',
+        hf_dataset_id='shibing624/AdvertiseGen',
+        preprocess_func=TextGenerationPreprocessor(
+            prompt=advertise_gen_prompt, columns={
+                'content': 'query',
+                'summary': 'response'
+            }),
+        tags=['text-generation', '🔥'],
+        split=['train', 'validation'],
+    ))
+
+
+class FireflyPreprocessor(ResponsePreprocessor):
+    _firefly_kind_list = {
+        'ProseGeneration', 'MRC', 'JinYongGeneration', 'TextCorrection', 'ClassicalChinese', 'BELLE', 'StoryGeneration',
+        'Couplet', 'Cot', 'Dictionary', 'Translation', 'Program', 'SentimentAnalyze', 'OpenQA', 'AncientPoem',
+        'TextMatching', 'NLI', 'Summary', 'KeywordRecognition', 'ProductDesc', 'LyricGeneration', 'Composition',
+        'MusicComment', 'NER'
+    }
+
+    def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        if row['kind'] not in FireflyPreprocessor._firefly_kind_list:
+            return
+        return super().preprocess(row)
+
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='AI-ModelScope/firefly-train-1.1M',
+        hf_dataset_id='YeungNLP/firefly-train-1.1M',
+        preprocess_func=FireflyPreprocessor(),
+        tags=['chat', 'general'],
+    ))
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='modelscope/clue',
+        hf_dataset_id='clue',
+        subsets=['cmnli'],
+        preprocess_func=ClsGenerationPreprocessor(['neutral', 'entailment', 'contradiction'],
+                                                  task='Natural Language Inference',
+                                                  is_pair_seq=True),
+        tags=['text-generation', 'classification'],
+        split=['train', 'validation'],
+    ))
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='DAMO_NLP/jd',
+        subsets=[
+            SubsetDataset(
+                'default',
+                'default',
+                preprocess_func=ClsGenerationPreprocessor(['negative', 'positive'],
+                                                          task='Sentiment Classification',
+                                                          is_pair_seq=False)),
+            SubsetDataset(
+                'cls',
+                'default',
+                preprocess_func=ClsPreprocessor(columns={'sentence': 'query'}),
+            ),
+        ],
+        tags=['text-generation', 'classification', '🔥'],
+        split=['train', 'validation'],
+    ))
+
+
+class SyntheticText2SqlPreprocessor(ResponsePreprocessor):
+
+    def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]:
+        sql_prompt = row['sql_prompt']
+        sql_context = row['sql_context']
+        sql = row['sql']
+        sql_explanation = row['sql_explanation']
+        query = f'Sql Table information:\n{sql_context}\n{sql_prompt}'
+        response = f'Let\'s think step by step:\n{sql_explanation}\nSo the final sql is:\n{sql}'
+        return super().preprocess({'query': query, 'response': response})
+
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='AI-ModelScope/synthetic_text_to_sql',
+        hf_dataset_id='gretelai/synthetic_text_to_sql',
+        preprocess_func=SyntheticText2SqlPreprocessor(),
+        tags=['nl2sql', 'en']))
+
+
+def _repair_toolbench(conversations: List[Dict[str, str]]) -> List[Dict[str, str]]:
+    assert len(conversations) == 2
+    if conversations[1]['from'] in {'caller', 'conclusion'}:
+        conversations[1]['from'] = 'assistant'
+    return conversations
+
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='shenweizhou/alpha-umi-toolbench-processed-v2',
+        subsets=['backbone', 'caller', 'planner', 'summarizer'],
+        preprocess_func=MessagesPreprocessor(repair_messages=_repair_toolbench),
+        tags=['chat', 'agent', '🔥'],
+        huge_dataset=True))
+
+
+class BlossomMathPreprocessor(ResponsePreprocessor):
+
+    def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]:
+        output, answer = row['output'], row['answer']
+        return super().preprocess({'query': row['query'], 'response': f'{output}\n\nAnswer: {answer}'})
+
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='AI-ModelScope/blossom-math-v2',
+        hf_dataset_id='Azure99/blossom-math-v2',
+        preprocess_func=BlossomMathPreprocessor(),
+        tags=['chat', 'math', '🔥']))
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='AI-ModelScope/sql-create-context',
+        hf_dataset_id='b-mc2/sql-create-context',
+        preprocess_func=AlpacaPreprocessor(columns={
+            'question': 'instruction',
+            'context': 'input',
+            'answer': 'output'
+        }),
+        tags=['chat', 'sql', '🔥']))
+
+
+class TigerBotLawPreprocessor(ResponsePreprocessor):
+
+    def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]:
+        prompt = """{type}
+{title}
+"""
+        cur_prompt = prompt.format(type=row['type'], title=row['title'])
+        for i in range(1, 4):
+            chapter = row[f'chapter{i}']
+            if chapter is not None:
+                cur_prompt += f'{chapter}'
+        cur_prompt += f'{row["response"]}'
+        return super().preprocess({'response': cur_prompt})
+
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='AI-ModelScope/tigerbot-law-plugin',
+        hf_dataset_id='TigerResearch/tigerbot-law-plugin',
+        preprocess_func=TigerBotLawPreprocessor(),
+        tags=['text-generation', 'law', 'pretrained']))
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='codefuse-ai/CodeExercise-Python-27k',
+        preprocess_func=MessagesPreprocessor(columns={'chat_rounds': 'messages'}),
+        tags=['chat', 'coding', '🔥']))
+
+
+class LeetcodePythonPreprocessor(ResponsePreprocessor):
+
+    def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]:
+        code_with_problem = row['code_with_problem']
+        idx = code_with_problem.find('```python')
+        problem = code_with_problem[:idx]
+        if problem.startswith('# '):
+            problem = problem[2:]
+        code = code_with_problem[idx:].strip()
+        explanation = row['explanation_only']
+        return super().preprocess({'query': problem, 'response': f'{code}\n\n{explanation}'})
+
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='AI-ModelScope/leetcode-solutions-python',
+        preprocess_func=LeetcodePythonPreprocessor(),
+        tags=['chat', 'coding', '🔥']))
+
+
+class StsbPreprocessor(RowPreprocessor):
+
+    def __init__(self, sim_threshold: Optional[float] = None):
+        self.sim_threshold = sim_threshold
+        super().__init__()
+
+    def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]:
+        row = {
+            'messages': [{
+                'role': 'user',
+                'content': row['sentence1']
+            }],
+            'positive_messages': [[{
+                'role': 'user',
+                'content': row['sentence2']
+            }]],
+            'label': row['score'],
+        }
+        if self.sim_threshold is None or float(row['label']) >= self.sim_threshold:
+            return row
+        else:
+            return None
+
+
+class StsbGeneratePreprocessor(ResponsePreprocessor):
+    prompt = """Task: Based on the given two sentences, provide a similarity score between 0.0 and 1.0.
+Sentence 1: {text1}
+Sentence 2: {text2}
+Similarity score: """
+
+    def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        return super().preprocess({
+            'query': self.prompt.format(text1=row['sentence1'], text2=row['sentence2']),
+            'response': f"{row['score']:.1f}"
+        })
+
+
+class StsbRegressionPreprocessor(StsbGeneratePreprocessor):
+
+    def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        return super(StsbGeneratePreprocessor, self).preprocess({
+            'query':
+            self.prompt.format(text1=row['sentence1'], text2=row['sentence2']),
+            'label':
+            row['score']
+        })
+
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='sentence-transformers/stsb',
+        hf_dataset_id='sentence-transformers/stsb',
+        subsets=[
+            SubsetDataset('default', preprocess_func=StsbPreprocessor()),  # embedding
+            SubsetDataset('positive', preprocess_func=StsbPreprocessor(sim_threshold=0.75)),  # infonce
+            SubsetDataset('generate', preprocess_func=StsbGeneratePreprocessor()),
+            SubsetDataset('reg', preprocess_func=StsbRegressionPreprocessor()),
+        ],
+        tags=['similarity', '🔥']))
+
+
+class MTEBRerankPreprocessor(RowPreprocessor):
+
+    def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]:
+        query = row['query']
+        positives = row['positive'] if isinstance(row['positive'], list) else [row['positive']]
+        negatives = row['negative'] if isinstance(row['negative'], list) else [row['negative']]
+
+        messages = [{'role': 'user', 'content': query}]
+        positive_messages = [[{'role': 'assistant', 'content': positive}] for positive in positives]
+        negative_messages = [[{'role': 'assistant', 'content': negative}] for negative in negatives]
+
+        return {'messages': messages, 'positive_messages': positive_messages, 'negative_messages': negative_messages}
+
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='MTEB/scidocs-reranking',
+        hf_dataset_id='mteb/scidocs-reranking',
+        split=['validation', 'test'],
+        preprocess_func=MTEBRerankPreprocessor(),
+        tags=['rerank', '🔥']))
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='MTEB/stackoverflowdupquestions-reranking',
+        hf_dataset_id='mteb/stackoverflowdupquestions-reranking',
+        split=['train', 'test'],
+        preprocess_func=MTEBRerankPreprocessor(),
+        tags=['rerank', '🔥']))
+
+
+def _repair_conversations_agent_instruct(s: str) -> List[Dict[str, Any]]:
+    s = s.replace('}\n {', '},\n {')
+    if isinstance(s, str):
+        s = ast.literal_eval(s)
+    return s
+
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='huangjintao/AgentInstruct_copy',
+        subsets=['alfworld', 'db', 'kg', 'mind2web', 'os', 'webshop'],
+        preprocess_func=MessagesPreprocessor(repair_messages=_repair_conversations_agent_instruct),
+        tags=['chat', 'agent', 'multi-round']))
+
+
+class MultiRoleAgentPreprocessor(RowPreprocessor):
+
+    def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        conv = row['conversations']
+        res_prompt = '\n\n【注意事项】\n1. 这是聊天室，不要发送私信给任何人\n2. 仅代表你个人说话,不要扮演其他人，只根据对话历史进行回复\n3. 长话短说，不要说太多话，不要超过50字 '
+        history_prompt = '\n\n【chat history】'
+        conv_prompt = '\n {name}:{content}'
+        query, response = '', conv[-1]['value']
+        system = conv[0]['value'] if conv[0]['from'] == 'system' else ''
+        if conv[0]['from'] == 'user':
+            query = conv[0]['value']
+        elif 'next_speakers:' not in system:
+            if '【注意事项】' not in system and system:
+                system += res_prompt
+            system += history_prompt
+            system += ''.join([conv_prompt.format(name=c['from'], content=c['value']) for c in conv[1:-1]])
+
+        if not query or not response:
+            return
+
+        return {
+            'messages': [{
+                'role': 'system',
+                'content': system
+            }, {
+                'role': 'user',
+                'content': query
+            }, {
+                'role': 'assistant',
+                'content': response
+            }],
+        }
+
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='iic/MSAgent-MultiRole',
+        preprocess_func=MultiRoleAgentPreprocessor(),
+        tags=['chat', 'agent', 'multi-round', 'role-play', 'multi-agent']))
+
+register_dataset(DatasetMeta(ms_dataset_id='swift/ToolBench', tags=['chat', 'agent', 'multi-round']))
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='tastelikefeet/competition_math',
+        subsets=[
+            SubsetDataset(
+                name='default',
+                subset='default',
+                split=['train', 'test'],
+            ),
+        ],
+        tags=['qa', 'math']))
+
+register_dataset(DatasetMeta(ms_dataset_id='modelscope/gsm8k', subsets=['main'], split=['train'], tags=['qa', 'math']))
+
+register_dataset(
+    DatasetMeta(ms_dataset_id='modelscope/MathR', subsets=['default', 'clean'], split=['train'], tags=['qa', 'math']))
+
+register_dataset(
+    DatasetMeta(ms_dataset_id='modelscope/MathR-32B-Distill', subsets=['data'], split=['train'], tags=['qa', 'math']))
+
+
+class CoundownTaskPreprocessor(ResponsePreprocessor):
+
+    def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]:
+        numbers = row['nums']
+        target = row.pop('response', None)
+        query = (f'Using the numbers {numbers}, create an equation that equals {target}.\n'
+                 'You can use basic arithmetic operations (+, -, *, /) and each number can only be used once.\n'
+                 'Show your work in <think> </think> tags. And return the final equation and answer '
+                 'in <answer> </answer> tags, for example <answer> (1 + 2) / 3 * 4 = 4 </answer>.')
+        row.update({'target': target, 'query': query})
+        return super().preprocess(row)
+
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='zouxuhong/Countdown-Tasks-3to4',
+        subsets=['default'],
+        preprocess_func=CoundownTaskPreprocessor(),
+        tags=['math']))
+
+
+class HC3Preprocessor(ResponsePreprocessor):
+    prompt = """Classification Task: Are the following responses from a human or from ChatGPT?
+Question: {question}
+Answer: {answer}
+Category: Human, ChatGPT
+Output:"""
+
+    def preprocess(self, row):
+        rows = []
+        for response in ['Human', 'ChatGPT']:
+            query = self.prompt.format(
+                question=row['query'], answer=self.random_state.choice(row[f'{response.lower()}_answers']))
+            rows.append(super().preprocess({'query': query, 'response': response}))
+        return rows
+
+
+class HC3ClsPreprocessor(HC3Preprocessor):
+
+    def preprocess(self, row):
+        rows = []
+        for i, response in enumerate(['Human', 'ChatGPT']):
+            query = self.prompt.format(
+                question=row['query'], answer=self.random_state.choice(row[f'{response.lower()}_answers']))
+            rows.append(ResponsePreprocessor.preprocess(self, {'query': query, 'label': i}))
+        return rows
+
+
+hc3_subset_names = ['baike', 'open_qa', 'nlpcc_dbqa', 'finance', 'medicine', 'law', 'psychology']
+hc3_subsets: List[SubsetDataset] = []
+for hc3_subset_name in hc3_subset_names:
+    hc3_subsets.append(
+        SubsetDataset(
+            name=hc3_subset_name,
+            subset=hc3_subset_name,
+            preprocess_func=HC3Preprocessor(),
+        ))
+    hc3_subsets.append(
+        SubsetDataset(
+            name=f'{hc3_subset_name}_cls',
+            subset=hc3_subset_name,
+            preprocess_func=HC3ClsPreprocessor(),
+        ))
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='simpleai/HC3-Chinese',
+        hf_dataset_id='Hello-SimpleAI/HC3-Chinese',
+        subsets=hc3_subsets,
+        tags=['text-generation', 'classification', '🔥']))
+
+hc3_subset_names = ['finance', 'medicine']
+hc3_subsets: List[SubsetDataset] = []
+for hc3_subset_name in hc3_subset_names:
+    hc3_subsets.append(
+        SubsetDataset(
+            name=hc3_subset_name,
+            subset=hc3_subset_name,
+            preprocess_func=HC3Preprocessor(),
+        ))
+    hc3_subsets.append(
+        SubsetDataset(
+            name=f'{hc3_subset_name}_cls',
+            subset=hc3_subset_name,
+            preprocess_func=HC3ClsPreprocessor(),
+        ))
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='simpleai/HC3',
+        hf_dataset_id='Hello-SimpleAI/HC3',
+        subsets=hc3_subsets,
+        preprocess_func=HC3Preprocessor(),
+        tags=['text-generation', 'classification', '🔥']))
+
+
+class DureaderPreprocessor(RowPreprocessor):
+
+    def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]:
+        prompt = """Task: Question Generation
+Context: {context}
+Answer: {answer}
+Question:"""
+        answer, context = row['text1'].split('[SEP]')
+        return {
+            'messages': [{
+                'role': 'user',
+                'content': prompt.format(context=context, answer=answer)
+            }, {
+                'role': 'assistant',
+                'content': row['text2']
+            }]
+        }
+
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='modelscope/DuReader_robust-QG',
+        preprocess_func=DureaderPreprocessor(),
+        split=['train', 'validation', 'test'],
+        tags=['text-generation', '🔥']))
+
+
+class HHRLHFPreprocessor(RowPreprocessor):
+
+    @staticmethod
+    def _to_messages(data):
+        messages = []
+        for query, response in zip(data[::2], data[1::2]):
+            messages.append({'role': 'user', 'content': query})
+            messages.append({'role': 'assistant', 'content': response})
+        return messages
+
+    def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]:
+        chosen = row['chosen'].strip()
+        rejected = row['rejected'].strip()
+        parts_chosen = [s.strip() for s in re.split('\n\nHuman:|\n\nAssistant:|\n\nHum:', chosen)]
+        parts_rejected = [s.strip() for s in re.split('\n\nHuman:|\n\nAssistant:|\n\nHum:', rejected)]
+        if parts_chosen[0].startswith('Human:'):
+            assert parts_rejected[0].startswith('Human:')
+            parts_chosen[0] = parts_chosen[0][6:].strip()
+            parts_rejected[0] = parts_rejected[0][6:].strip()
+        row['messages'] = self._to_messages(parts_chosen)
+        row['rejected_messages'] = self._to_messages(parts_rejected)
+        return row
+
+
+# TODO meta file broken
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='AI-ModelScope/hh-rlhf',
+        subsets=['helpful-base', 'helpful-online', 'helpful-rejection-sampled'],
+        preprocess_func=HHRLHFPreprocessor(),
+        split=['train', 'test'],
+        tags=['rlhf', 'dpo'],
+        huge_dataset=True))
+
+
+class XlamFunctionCallingPreprocessor(RowPreprocessor):
+
+    def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]:
+        messages = [{'role': 'user', 'content': row['query']}]
+        response = row['answers']
+        response = json.loads(response)
+        messages += [{'role': 'tool_call', 'content': json.dumps(content)} for content in response]
+        return {'messages': messages, 'tools': row['tools']}
+
+
+class XlamFunctionCallingGRPOPreprocessor(ResponsePreprocessor):
+
+    def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]:
+        query = row['query']
+        answers = row['response']
+        if isinstance(answers, str):
+            answers = json.loads(answers)
+        answer = np.random.choice(answers)
+        name = answer['name']
+        args = json.dumps(answer['arguments'])
+        response = f'Action: {name}\nAction Input: {args}'
+        row = {'query': query, 'response': response, 'solution': response, 'tools': row['tools']}
+        return super().preprocess(row)
+
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='LLM-Research/xlam-function-calling-60k',
+        hf_dataset_id='Salesforce/xlam-function-calling-60k',
+        subsets=[
+            SubsetDataset('default', 'dataset', preprocess_func=XlamFunctionCallingPreprocessor()),
+            SubsetDataset('grpo', 'dataset', preprocess_func=XlamFunctionCallingGRPOPreprocessor())
+        ],
+        tags=['agent', 'grpo', '🔥']))
+
+
+class HHRLHFCNPreprocessor(MessagesPreprocessor):
+
+    def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]:
+        row['messages'].append(row.pop('chosen'))
+        row['rejected_response'] = row['rejected']['text']
+        return super().preprocess(row)
+
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='AI-ModelScope/hh_rlhf_cn',
+        subsets=['hh_rlhf', 'harmless_base_cn', 'harmless_base_en', 'helpful_base_cn', 'helpful_base_en'],
+        preprocess_func=HHRLHFCNPreprocessor(columns={'context': 'messages'}, content_key='text'),
+        split=['train', 'test'],
+        tags=['rlhf', 'dpo', '🔥']))
+
+
+def repair_conversations(s: Union[str, Any]) -> Any:
+    if isinstance(s, str):
+        s = s.replace('}\n {', '},{')
+        s = s.replace('}\n{', '},{')
+        s = s.replace('}{', '},{')
+        s = s.replace('}\n  {', '},{')
+        return ast.literal_eval(s)
+    return s
+
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='AI-ModelScope/lmsys-chat-1m',
+        hf_dataset_id='lmsys/lmsys-chat-1m',
+        preprocess_func=MessagesPreprocessor(repair_messages=repair_conversations),
+        tags=['chat', 'em']))
+
+
+class EmojiPreprocessr(ResponsePreprocessor):
+
+    def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]:
+        # Remove dirty characters
+        row['query'] = row['query'].replace('️', '')
+        row['response'] = row['response'].replace('️', '')
+        row['rejected_response'] = row['rejected_response'].replace('️', '')
+        return super().preprocess(row)
+
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='hjh0119/shareAI-Llama3-DPO-zh-en-emoji',
+        hf_dataset_id='shareAI/DPO-zh-en-emoji',
+        preprocess_func=EmojiPreprocessr(columns={
+            'answer_zh': 'response',
+            'answer_en': 'rejected_response'
+        }),
+        tags=['rlhf', 'dpo']))
+
+register_dataset(
+    DatasetMeta(ms_dataset_id='AI-ModelScope/ultrafeedback-binarized-preferences-cleaned-kto', tags=['rlhf', 'kto']))
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='OmniData/Zhihu-KOL-More-Than-100-Upvotes',
+        hf_dataset_id='bzb2023/Zhihu-KOL-More-Than-100-Upvotes',
+        tags=['zhihu', 'qa']))
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='OmniData/Zhihu-KOL',
+        hf_dataset_id='wangrui6/Zhihu-KOL',
+        huge_dataset=True,
+        tags=['zhihu', 'qa'],
+    ))
+
+
+class GuanacoPreprocessor(RowPreprocessor):
+
+    def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        instruction = row['instruction']
+        input = row['input']
+        output = row['output']
+        history = []
+        if instruction:
+            parts = split_str_parts_by(
+                instruction, ['User:', 'User：', 'Assistant：', 'Assistant:', 'Asssistent:', 'Assistent:', 'Assistenz:'])
+            for idx, part in enumerate(parts):
+                if idx % 2 == 0:
+                    if 'user' not in part['key'].lower():
+                        return
+                    history.append([part['content'], None])
+                else:
+                    if 'assist' not in part['key'].lower() and 'asssist' not in part['key'].lower():
+                        return
+                    history[-1][-1] = part['content']
+        if input.startswith('User:'):
+            input = input[len('User:'):].strip()
+        if any([not h[0] or not h[1] for h in history]):
+            return
+
+        messages = []
+        for h in history:
+            messages.append({'role': 'user', 'content': h[0]})
+            messages.append({'role': 'assistant', 'content': h[1]})
+        messages.append({'role': 'user', 'content': input})
+        messages.append({'role': 'assistant', 'content': output})
+        return {
+            'messages': messages,
+        }
+
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='AI-ModelScope/GuanacoDataset',
+        hf_dataset_id='JosephusCheung/GuanacoDataset',
+        preprocess_func=GuanacoPreprocessor(),
+        tags=['chat', 'zh']))
+
+
+class FunctionCallChatmlPreprocessor(MessagesPreprocessor):
+
+    def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        res = super().preprocess(row)
+
+        if res['function_description']:
+            res['tools'] = res['function_description'].split('\n\n')
+        messages = res['messages']
+        if messages[0]['role'] == 'system':
+            messages.pop(0)
+        return res
+
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='AI-ModelScope/function-calling-chatml',
+        hf_dataset_id='Locutusque/function-calling-chatml',
+        preprocess_func=FunctionCallChatmlPreprocessor(),
+        tags=['agent', 'en', 'sft', '🔥']))
+
+
+class Dolly15kPreprocessor(RowPreprocessor):
+
+    def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        instruction = row['instruction']
+        context = row['context']
+        response = row['response']
+        query = ''
+        if context:
+            query = 'Here gives some useful information:\n'
+            query += context
+            query += '\n'
+        query += instruction
+        return {
+            'messages': [{
+                'role': 'user',
+                'content': query
+            }, {
+                'role': 'assistant',
+                'content': response
+            }],
+        }
+
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='AI-ModelScope/databricks-dolly-15k',
+        hf_dataset_id='databricks/databricks-dolly-15k',
+        preprocess_func=Dolly15kPreprocessor(),
+        tags=['multi-task', 'en', 'quality']))
+
+
+class OrpoDPOMix40kPreprocessor(MessagesPreprocessor):
+
+    def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        if row['source'] == 'toxic-dpo-v0.2':
+            return
+        return super().preprocess(row)
+
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='AI-ModelScope/orpo-dpo-mix-40k',
+        hf_dataset_id='mlabonne/orpo-dpo-mix-40k',
+        preprocess_func=OrpoDPOMix40kPreprocessor(columns={
+            'chosen': 'messages',
+            'rejected': 'rejected_messages'
+        }),
+        tags=['dpo', 'orpo', 'en', 'quality']))
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='swift/sharegpt',
+        subsets=['common-zh', 'unknow-zh', 'common-en'],
+        tags=['chat', 'general', 'multi-round']))
+
+
+class SelfCognitionPreprocessor(ResponsePreprocessor):
+
+    def __init__(self, *args, query_suffix: str = '', response_prefix: str = '', **kwargs):
+        self.query_suffix = query_suffix
+        self.response_prefix = response_prefix
+        self.name: Optional[Tuple[str, str]] = None
+        self.author: Optional[Tuple[str, str]] = None
+        super().__init__(*args, **kwargs)
+
+    def set_name_author(self, name, author):
+        self.name = name
+        self.author = author
+
+    def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]:
+        for key in ['name', 'author']:
+            val = getattr(self, key)
+            if val is None:
+                continue
+            val = val[0] if row['tag'] == 'zh' else val[1]
+            if val is None:
+                continue
+            placeholder = '{{' + key.upper() + '}}'
+            row['query'] = row['query'].replace(placeholder, val)
+            row['response'] = row['response'].replace(placeholder, val)
+
+        row['query'] = row['query'] + self.query_suffix
+        row['response'] = self.response_prefix + row['response']
+        return super().preprocess(row)
+
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='swift/self-cognition',
+        hf_dataset_id='modelscope/self-cognition',
+        subsets=[
+            SubsetDataset(preprocess_func=SelfCognitionPreprocessor()),
+            SubsetDataset(
+                'qwen3',
+                preprocess_func=SelfCognitionPreprocessor(
+                    query_suffix=' /no_think', response_prefix='<think>\n\n</think>\n\n')),
+            SubsetDataset(
+                'empty_think', preprocess_func=SelfCognitionPreprocessor(response_prefix='<think>\n\n</think>\n\n')),
+        ],
+        dataset_name='self-cognition',
+        tags=['chat', 'self-cognition', '🔥']))
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='open-r1/DAPO-Math-17k-Processed',
+        hf_dataset_id='open-r1/DAPO-Math-17k-Processed',
+        subsets=['all'],
+        tags=['math', 'rlvr']))
+
+
+class SudokuPreprocessor(ResponsePreprocessor):
+    prompt = ('Solve the following 9x9 Sudoku puzzle. '
+              "Empty cells are marked with '0'. "
+              'Provide the completed grid as your answer.\n\n'
+              'Puzzle:\n{puzzle}')
+
+    @staticmethod
+    def _format_grid(s: str) -> str:
+        return '\n'.join(s[i:i + 9] for i in range(0, len(s), 9))
+
+    def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        puzzle = row['query'].replace('.', '0')
+        response = row['response']
+        puzzle = self._format_grid(puzzle)
+        response = self._format_grid(response)
+        return super().preprocess({'query': self.prompt.format(puzzle=puzzle), 'response': response})
+
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='sapientinc/sudoku-extreme-1k',
+        hf_dataset_id='sapientinc/sudoku-extreme-1k',
+        preprocess_func=SudokuPreprocessor(),
+    ))

swift/dataset/dataset_meta.py

@@ -0,0 +1,202 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import numpy as np
+import os
+import shutil
+from abc import ABC, abstractmethod
+from copy import deepcopy
+from dataclasses import dataclass, field
+from datasets import Dataset as HfDataset
+from datasets import concatenate_datasets, interleave_datasets
+from modelscope.hub.api import ModelScopeConfig
+from modelscope.utils.config_ds import MS_CACHE_HOME
+from tempfile import TemporaryDirectory
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, Union
+
+from swift.utils import download_ms_file, get_logger, get_seed, safe_ddp_context
+from .preprocessor import DATASET_TYPE, AutoPreprocessor
+from .utils import sample_dataset
+
+PreprocessFunc = Callable[..., DATASET_TYPE]
+logger = get_logger()
+
+if TYPE_CHECKING:
+    from .dataset_syntax import DatasetSyntax
+
+
+@dataclass
+class SubsetDataset:
+    # `Name` is used for matching subsets of the dataset, and `subset` refers to the subset_name on the hub.
+    name: Optional[str] = None
+    # If set to None, then subset is set to subset_name.
+    subset: str = 'default'
+
+    # Higher priority. If set to None, the attributes of the DatasetMeta will be used.
+    split: Optional[List[str]] = None
+    preprocess_func: Optional[PreprocessFunc] = None
+
+    # If the dataset specifies "all," weak subsets will be skipped.
+    is_weak_subset: bool = False
+
+    def __post_init__(self):
+        if self.name is None:
+            self.name = self.subset
+
+    def set_default(self, dataset_meta: 'DatasetMeta') -> 'SubsetDataset':
+        subset_dataset = deepcopy(self)
+        for k in ['split', 'preprocess_func']:
+            v = getattr(subset_dataset, k)
+            if v is None:
+                setattr(subset_dataset, k, deepcopy(getattr(dataset_meta, k)))
+        return subset_dataset
+
+
+class BaseDatasetLoader(ABC):
+
+    @abstractmethod
+    def load(
+        self,
+        dataset_syntax: Optional['DatasetSyntax'] = None,
+        dataset_meta: Optional['DatasetMeta'] = None,
+        *,
+        use_hf: Optional[bool] = None,
+    ) -> HfDataset:
+        pass
+
+    @staticmethod
+    def download_ms_dataset(ms_dataset_id: str, files: List[str], force_download: bool = False) -> str:
+        """Download dataset from repo manually
+        Args:
+            ms_dataset_id: The dataset id of ModelScope
+            files: Which files to download
+            force_download: Force download or not
+        Returns:
+            The dataset dir
+        """
+        assert isinstance(files, list)
+        url = f'http://www.modelscope.cn/api/v1/datasets/{ms_dataset_id}/repo?Revision=master&FilePath={{fpath}}'
+        cache_dir = os.path.join(MS_CACHE_HOME, 'datasets', ms_dataset_id, 'master')
+        local_dir = os.path.join(cache_dir, 'raw')
+        tmp_dir = os.path.join(cache_dir, 'tmp')
+        os.makedirs(local_dir, exist_ok=True)
+        os.makedirs(tmp_dir, exist_ok=True)
+        cookies = ModelScopeConfig.get_cookies()
+        with TemporaryDirectory(dir=tmp_dir) as temp_dir:
+            for remote_fpath in files:
+                url = url.format(fpath=remote_fpath)
+                temp_fpath = os.path.join(temp_dir, remote_fpath)
+                local_fpath = os.path.join(local_dir, remote_fpath)
+                if not force_download and os.path.exists(local_fpath):
+                    continue
+                download_ms_file(url, temp_fpath, cookies)
+                shutil.copy2(temp_fpath, local_fpath)
+
+        return local_dir
+
+    @staticmethod
+    def concat_datasets(datasets: List[HfDataset]) -> Optional[HfDataset]:
+        if len(datasets) == 0:
+            return
+        if len(datasets) == 1:
+            return datasets[0]
+        return concatenate_datasets(datasets)
+
+    @staticmethod
+    def interleave_datasets(datasets, *args, **kwargs):
+        if len(datasets) == 0:
+            return
+        if len(datasets) == 1:
+            return datasets[0]
+        return interleave_datasets(datasets, *args, **kwargs)
+
+    @staticmethod
+    def shuffle_dataset(dataset, seed: int, buffer_size: int = 1000):
+        if isinstance(dataset, HfDataset):
+            with safe_ddp_context(None, True):
+                return dataset.shuffle(seed=seed)
+        else:
+            return dataset.shuffle(seed=seed, buffer_size=buffer_size)
+
+    @staticmethod
+    def post_process(
+        train_dataset: DATASET_TYPE,
+        *,
+        dataset_sample: Optional[int] = None,
+        split_dataset_ratio: float = 0.,
+        streaming: bool = False,
+        shuffle: bool = True,
+        random_state: Optional[np.random.RandomState] = None,
+    ) -> Tuple[DATASET_TYPE, Optional[DATASET_TYPE]]:
+        """Split into train/val datasets and perform dataset sampling."""
+        assert dataset_sample is None or dataset_sample > 0
+        assert 0 <= split_dataset_ratio <= 1
+        if streaming:
+            if dataset_sample is None:
+                if split_dataset_ratio == 0:
+                    val_dataset = None
+                elif split_dataset_ratio == 1:
+                    train_dataset, val_dataset = None, train_dataset
+                else:
+                    raise ValueError('The IterableDataset does not support splitting the training set '
+                                     'and validation set when dataset_sample is None.')
+            else:
+                # not shuffle
+                train_dataset = train_dataset.take(dataset_sample)
+                val_sample = int(dataset_sample * split_dataset_ratio)
+                val_dataset = None if val_sample == 0 else train_dataset.take(val_sample)
+                if val_sample:
+                    train_dataset = train_dataset.skip(val_sample)
+        else:
+            if dataset_sample is None:
+                dataset_sample = len(train_dataset)
+            if split_dataset_ratio == 0:
+                train_dataset = sample_dataset(train_dataset, dataset_sample, shuffle, random_state)
+                val_dataset = None
+            elif split_dataset_ratio == 1:
+                train_dataset, val_dataset = None, train_dataset
+                val_sample = dataset_sample
+                # Avoid duplication in the val_dataset.
+                assert val_sample <= len(val_dataset), f'val_sample: {val_sample}, len(val_dataset): {len(val_dataset)}'
+                val_dataset = sample_dataset(val_dataset, val_sample, shuffle, random_state)
+            else:
+                # Avoid duplication in the val_dataset.
+                train_len = min(len(train_dataset), dataset_sample)
+                val_sample = max(int(train_len * split_dataset_ratio), 1)
+                train_sample = dataset_sample - val_sample
+                assert train_sample > 0
+                with safe_ddp_context(None, True):
+                    train_dataset, val_dataset = train_dataset.train_test_split(
+                        test_size=val_sample, shuffle=shuffle, seed=get_seed(random_state)).values()
+                train_dataset = sample_dataset(train_dataset, train_sample, shuffle, random_state)
+        return train_dataset, val_dataset
+
+
+@dataclass
+class DatasetMeta:
+    ms_dataset_id: Optional[str] = None
+    hf_dataset_id: Optional[str] = None
+    dataset_path: Optional[str] = None  # or dataset_dir
+    dataset_name: Optional[str] = None
+    ms_revision: Optional[str] = None
+    hf_revision: Optional[str] = None
+
+    subsets: List[Union[SubsetDataset, str]] = field(default_factory=lambda: ['default'])
+    # Applicable to all subsets.
+    split: List[str] = field(default_factory=lambda: ['train'])
+    # First perform column mapping, then proceed with the preprocess_func.
+    preprocess_func: PreprocessFunc = field(default_factory=lambda: AutoPreprocessor())
+    loader: Optional[BaseDatasetLoader] = None
+
+    tags: List[str] = field(default_factory=list)
+    help: Optional[str] = None
+    huge_dataset: bool = False
+
+    def __post_init__(self):
+        from .loader import DatasetLoader
+        if self.loader is None:
+            self.loader = DatasetLoader
+        for i, subset in enumerate(self.subsets):
+            if isinstance(subset, str):
+                self.subsets[i] = SubsetDataset(subset=subset)
+
+
+DATASET_MAPPING: Dict[Tuple[str, str, str], DatasetMeta] = {}

swift/dataset/dataset_syntax.py

@@ -0,0 +1,127 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import os
+import platform
+import re
+from dataclasses import dataclass, field
+from typing import Dict, List, Literal, Optional, Tuple
+
+from .dataset_meta import DATASET_MAPPING, DatasetMeta
+
+_dataset_meta_mapping = None
+
+
+@dataclass
+class DatasetSyntax:
+    dataset: str
+    subsets: List[str] = field(default_factory=list)
+    dataset_sample: Optional[int] = None
+    use_hf: Optional[bool] = None
+
+    def __post_init__(self):
+        if os.path.isfile(self.dataset):
+            self.dataset_type = 'path'
+        else:  # dataset_id or dataset_dir
+            self.dataset_type = 'repo'
+
+    def get_raw(self):
+        subsets = '/'.join(self.subsets)
+        dataset_sample = '' if self.dataset_sample is None else f'#{self.dataset_sample}'
+        return f'{self.dataset}{subsets}{dataset_sample}'
+
+    @staticmethod
+    def _safe_split(s: str,
+                    sep: str,
+                    use_0: bool,
+                    split_mode: Literal['left', 'right'] = 'left') -> Tuple[Optional[str], Optional[str]]:
+        """
+        use_0: When the length of the part is 1, is it considered as part0 or part1.
+        split_mode: use split or rsplit
+        """
+        if s is None or len(s) == 0:
+            return None, None
+        if split_mode == 'left':
+            part = s.split(sep, 1)
+        else:
+            part = s.rsplit(sep, 1)
+        if len(part) == 1:
+            if use_0:
+                part = part[0], None
+            else:
+                part = None, part[0]
+        else:
+            assert len(part) == 2
+        return part
+
+    @classmethod
+    def parse(cls, dataset: str) -> 'DatasetSyntax':
+        """Parse the dataset from the command line"""
+        # hf/ms::dataset_id or dataset_path:subset1/subset2/subset3#dataset_sample
+        if os.path.exists(dataset):
+            use_hf = None
+        else:
+            use_hf, dataset = cls._safe_split(dataset, '::', False)
+            if isinstance(use_hf, str):
+                use_hf = use_hf.lower()
+            use_hf = {'hf': True, 'ms': False}.get(use_hf)
+        if os.path.exists(dataset):
+            other, dataset_sample = dataset, None
+        else:
+            other, dataset_sample = cls._safe_split(dataset, '#', True, 'right')
+        if os.path.exists(other):
+            dataset, subsets = other, None
+        else:
+            dataset, subsets = cls._safe_split(other, ':', True)
+
+        if subsets is not None:
+            subsets = [subset.strip() for subset in subsets.split('/')]
+        if dataset_sample is not None:
+            dataset_sample = int(dataset_sample)
+        return cls(dataset.strip(), subsets or [], dataset_sample, use_hf)
+
+    def get_dataset_meta(self, use_hf: bool):
+        dataset_meta_mapping = self._get_dataset_meta_mapping()
+        dataset_type = self.dataset_type
+        if dataset_type == 'path':
+            dataset_meta = dataset_meta_mapping.get((dataset_type, self.dataset))
+        else:
+            dataset_type = 'repo' if os.path.isdir(self.dataset) else {True: 'hf', False: 'ms'}[use_hf]
+            dataset_meta = dataset_meta_mapping.get((dataset_type, self.dataset))
+        return dataset_meta or self._get_matched_dataset_meta(dataset_meta_mapping) or DatasetMeta()
+
+    @staticmethod
+    def _get_dataset_meta_mapping() -> Dict[Tuple[str, str], DatasetMeta]:
+        global _dataset_meta_mapping
+        if _dataset_meta_mapping is not None:
+            return _dataset_meta_mapping
+        _dataset_meta_mapping = {}
+        for dataset_meta in DATASET_MAPPING.values():
+            if dataset_meta.dataset_path is not None:
+                dataset_type = 'repo' if os.path.isdir(dataset_meta.dataset_path) else 'path'
+                _dataset_meta_mapping[(dataset_type, dataset_meta.dataset_path)] = dataset_meta
+            if dataset_meta.ms_dataset_id is not None:
+                _dataset_meta_mapping[('ms', dataset_meta.ms_dataset_id)] = dataset_meta
+            if dataset_meta.hf_dataset_id is not None:
+                _dataset_meta_mapping[('hf', dataset_meta.hf_dataset_id)] = dataset_meta
+        return _dataset_meta_mapping
+
+    @staticmethod
+    def get_dataset_name(dataset_id: str) -> str:
+        # compat hf hub
+        dataset_id = dataset_id.rstrip('/')
+        match_ = re.search('/datasets--.+?--(.+?)/snapshots/', dataset_id)
+        if match_ is not None:
+            return match_.group(1)
+
+        dataset_name = dataset_id.rsplit('/', 1)[-1]
+        if platform.system().lower() == 'windows':
+            dataset_name = dataset_name.rsplit('\\', 1)[-1]
+        return dataset_name
+
+    def _get_matched_dataset_meta(self, dataset_meta_mapping):
+        suffix_dataset_meta_mapping = {}
+        for dataset_name, dataset_meta in dataset_meta_mapping.items():
+            dataset_name = self.get_dataset_name(dataset_name[1])
+            suffix_dataset_meta_mapping[dataset_name] = dataset_meta
+        dataset_name = self.get_dataset_name(self.dataset)
+        dataset_meta = suffix_dataset_meta_mapping.get(dataset_name)
+        return dataset_meta

swift/dataset/indexed_dataset.py

@@ -0,0 +1,132 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import bisect
+import mmap
+import os
+import pickle
+import threading
+from modelscope.hub.utils.utils import get_cache_dir
+from queue import Queue
+from torch.utils.data import Dataset
+from typing import Any, List
+
+
+class IndexedDatasetBuilder:
+    CHUNK_SIZE = 1e10
+
+    def __init__(self, dataset_name: str):
+        self.cache_dir = IndexedDataset.get_cache_dir(dataset_name)
+        self.n_shard = 1
+        self.bin_path = os.path.join(self.cache_dir, IndexedDataset.BIN_FNAME.format(0))
+        self.idx_path = os.path.join(self.cache_dir, IndexedDataset.IDX_FNAME)
+        if os.path.exists(self.bin_path):
+            os.remove(self.bin_path)
+        self.bin_file = open(self.bin_path, 'ab')
+        self.length_list = []
+        self.idx_list = [0]
+        self.shard_offset = [0]
+        self._thread = None
+        self._queue = Queue(maxsize=1000)
+
+    def _write_worker(self):
+        while True:
+            items = self._queue.get()
+            if items is None:
+                break
+            bin_buffer = []
+            for item in items:
+                item_buffer = pickle.dumps(item)
+                bin_buffer.append(item_buffer)
+                self.idx_list.append(self.idx_list[-1] + len(item_buffer))
+                self.length_list.append(item['length'])
+            self.bin_file.write(b''.join(bin_buffer))
+            offset = self.idx_list[-1] - self.shard_offset[-1]
+            if offset >= self.CHUNK_SIZE:
+                self.bin_file.close()
+                self.bin_path = os.path.join(self.cache_dir, IndexedDataset.BIN_FNAME.format(self.n_shard))
+                self.shard_offset.append(self.shard_offset[-1] + offset)
+                self.n_shard += 1
+                if os.path.exists(self.bin_path):
+                    os.remove(self.bin_path)
+                self.bin_file = open(self.bin_path, 'ab')
+
+    def add_items(self, items: List[Any]) -> None:
+        if self._thread is None:
+            self._thread = threading.Thread(target=self._write_worker, daemon=True)
+            self._thread.start()
+        self._queue.put(items)
+
+    def finalize(self):
+        if self._thread is not None:
+            self._queue.put(None)
+            self._thread.join()
+        self.bin_file.close()
+        idx_obj = {
+            'idx': self.idx_list,
+            'length': self.length_list,
+            'n_shard': self.n_shard,
+            'shard_offset': self.shard_offset,
+        }
+        with open(self.idx_path, 'wb') as f:
+            pickle.dump(idx_obj, f)
+
+
+class BinReader:
+
+    def __init__(self, bin_path: str):
+        self.bin_path = bin_path
+        self.file = open(bin_path, 'rb')
+        try:
+            self.mm = mmap.mmap(self.file.fileno(), 0, access=mmap.ACCESS_READ)
+        except ValueError:
+            # For example, self.file is an empty file.
+            self.mm = None
+
+    def read_buffer(self, offset: int, size: int) -> bytes:
+        if offset < 0 or size < 0 or offset + size > len(self.mm):
+            raise ValueError('Invalid offset or size')
+        return self.mm[offset:offset + size]
+
+    def __del__(self):
+        if self.mm is not None:
+            self.mm.close()
+        self.file.close()
+
+
+class IndexedDataset(Dataset):
+    BIN_FNAME = 'data-{:05d}.bin'
+    IDX_FNAME = 'data.idx'
+
+    @staticmethod
+    def get_cache_dir(dataset_name: str):
+        cache_dir = os.getenv('PACKING_CACHE') or os.path.join(get_cache_dir(), 'tmp')
+        cache_dir = os.path.join(cache_dir, dataset_name)
+        os.makedirs(cache_dir, exist_ok=True)
+        assert dataset_name is not None, f'dataset_name: {dataset_name}'
+        return cache_dir
+
+    def __init__(self, dataset_name: str):
+        self.dataset_name = dataset_name
+        cache_dir = self.get_cache_dir(dataset_name)
+        self.idx_path = os.path.join(cache_dir, self.IDX_FNAME)
+        with open(self.idx_path, 'rb') as f:
+            idx_obj = pickle.load(f)
+        self.idx_list = idx_obj['idx']
+        self.length_list = idx_obj['length']
+        self.n_shard = idx_obj['n_shard']
+        self.shard_offset = idx_obj['shard_offset']
+        self.bin_readers = []
+        for i in range(self.n_shard):
+            bin_path = os.path.join(cache_dir, self.BIN_FNAME.format(i))
+            self.bin_readers.append(BinReader(bin_path))
+
+    def __getitem__(self, index: int):
+        if index < 0:
+            index = index % len(self)
+        idx, idx_next = self.idx_list[index], self.idx_list[index + 1]
+        num_shard = bisect.bisect_right(self.shard_offset, idx)
+        offset = self.shard_offset[num_shard - 1]
+        buffer = self.bin_readers[num_shard - 1].read_buffer(idx - offset, idx_next - idx)
+        return pickle.loads(buffer)
+
+    def __len__(self):
+        return len(self.idx_list) - 1

swift/dataset/loader.py

@@ -0,0 +1,384 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import numpy as np
+import os
+from contextlib import nullcontext
+from datasets import Dataset as HfDataset
+from datasets import IterableDataset as HfIterableDataset
+from datasets import load_dataset as hf_load_dataset
+from functools import partial
+from modelscope.hub.utils.utils import get_cache_dir
+from typing import Dict, List, Literal, Optional, Tuple, Union
+
+from swift.hub import get_hub
+from swift.utils import get_logger, get_seed, safe_ddp_context, use_hf_hub
+from .dataset_meta import DATASET_TYPE, BaseDatasetLoader
+from .dataset_syntax import DatasetSyntax
+from .preprocessor import RowPreprocessor
+from .register import DATASET_MAPPING, DatasetMeta, SubsetDataset
+
+logger = get_logger()
+
+
+class DatasetLoader(BaseDatasetLoader):
+
+    def __init__(
+        self,
+        num_proc: int = 1,
+        load_from_cache_file: bool = True,
+        streaming: bool = False,
+        hub_token: Optional[str] = None,
+        strict: bool = False,
+        download_mode: Literal['force_redownload', 'reuse_dataset_if_exists'] = 'reuse_dataset_if_exists',
+        columns: Optional[Dict[str, str]] = None,
+        remove_unused_columns: bool = True,
+        disable_auto_column_mapping: bool = False,
+    ):
+        self.num_proc = num_proc
+        self.load_from_cache_file = load_from_cache_file
+        self.streaming = streaming
+        self.hub_token = hub_token
+        self.strict = strict
+        self.download_mode = download_mode
+        self.columns = columns
+        self.remove_unused_columns = remove_unused_columns
+        self.disable_auto_column_mapping = disable_auto_column_mapping
+
+    def _load_dataset_path(
+        self,
+        dataset_path: str,
+        dataset_meta: DatasetMeta,
+    ) -> HfDataset:
+        ext = os.path.splitext(dataset_path)[1].lstrip('.')
+        file_type = {'jsonl': 'json', 'txt': 'text'}.get(ext) or ext
+        kwargs = {'split': 'train', 'streaming': self.streaming, 'num_proc': self.num_proc}
+        if file_type == 'csv':
+            kwargs['na_filter'] = False
+        with safe_ddp_context(None, True):
+            kwargs['cache_dir'] = os.path.join(get_cache_dir(), 'datasets')
+            dataset = hf_load_dataset(file_type, data_files=dataset_path, **kwargs)
+        if self.columns:
+            dataset = RowPreprocessor.safe_rename_columns(dataset, self.columns)
+        dataset = dataset_meta.preprocess_func(
+            dataset,
+            num_proc=self.num_proc,
+            load_from_cache_file=self.load_from_cache_file,
+            strict=self.strict,
+            enable_auto_mapping=not self.disable_auto_column_mapping)
+        if self.remove_unused_columns:
+            dataset = RowPreprocessor.remove_useless_columns(dataset)
+        return dataset
+
+    def _load_repo_dataset(
+        self,
+        dataset_id: str,
+        subset: SubsetDataset,
+        *,
+        use_hf: Optional[bool] = None,
+        revision: Optional[str] = None,
+    ) -> HfDataset:
+        datasets = []
+        if os.path.isdir(dataset_id):
+            retry = 1
+            load_context = nullcontext
+            use_hf = True
+            dataset_str = f'Use local folder, dataset_dir: {dataset_id}'
+            # The dataset downloaded from modelscope will have an additional dataset_infos.json file.
+            with safe_ddp_context('dataset_infos_rename'):
+                dataset_infos_path = os.path.join(dataset_id, 'dataset_infos.json')
+                if os.path.isfile(dataset_infos_path):
+                    os.rename(dataset_infos_path, f'{dataset_infos_path}_bak')
+        elif dataset_id.startswith('/'):
+            raise ValueError(f'The local path does not exist, dataset_id: `{dataset_id}`. '
+                             f'os.path.exists(dataset_id): {os.path.exists(dataset_id)}')
+        else:
+            retry = 3
+            load_context = partial(safe_ddp_context, hash_id=dataset_id, use_barrier=True)
+            dataset_str_f = 'Downloading the dataset from {hub}, dataset_id: {dataset_id}'
+            if use_hf:
+                dataset_str = dataset_str_f.format(hub='HuggingFace', dataset_id=dataset_id)
+            else:
+                dataset_str = dataset_str_f.format(hub='ModelScope', dataset_id=dataset_id)
+        logger.info(dataset_str)
+        hub = get_hub(use_hf)
+        for split in subset.split:
+            i = 1
+            with load_context():
+                while True:
+                    try:
+                        dataset = hub.load_dataset(
+                            dataset_id,
+                            subset.subset,
+                            split,
+                            streaming=self.streaming,
+                            revision=revision,
+                            download_mode=self.download_mode,
+                            hub_token=self.hub_token,
+                            num_proc=self.num_proc)
+                    except Exception as e:
+                        if i == retry:
+                            raise
+                        i += 1
+                        logger.error(f'Dataset {dataset_id} load failed: subset_name={subset.subset},'
+                                     f'split={split} with error: {e}')
+                    else:
+                        break
+            if hasattr(dataset, '_hf_ds'):
+                dataset = dataset._hf_ds
+                if self.streaming and isinstance(dataset, HfDataset):
+                    dataset = dataset.to_iterable_dataset()
+            if self.columns:
+                dataset = RowPreprocessor.safe_rename_columns(dataset, self.columns)
+            dataset = subset.preprocess_func(
+                dataset,
+                num_proc=self.num_proc,
+                load_from_cache_file=self.load_from_cache_file,
+                strict=self.strict,
+                enable_auto_mapping=not self.disable_auto_column_mapping)
+            if self.remove_unused_columns:
+                dataset = RowPreprocessor.remove_useless_columns(dataset)
+            datasets.append(dataset)
+        return self.concat_datasets(datasets)
+
+    @staticmethod
+    def _select_subsets(subsets: List[str], dataset_meta: DatasetMeta) -> List[SubsetDataset]:
+        subset_mapping = {subset.name: subset for subset in dataset_meta.subsets}
+        subset_names = list(subset_mapping.keys())
+        if not subsets:
+            if len(subset_names) <= 1:
+                subsets = subset_names
+            elif 'default' in subset_names:
+                subsets = ['default']
+            else:
+                raise ValueError(f'Please provide subsets. available subsets: {subset_names}')
+        elif len(subsets) == 1 and subsets[0] == 'all' and 'all' not in subset_names:
+            subsets = [subset_name for subset_name in subset_names if not subset_mapping[subset_name].is_weak_subset]
+
+        subsets = [
+            subset_mapping[subset_name] if subset_name in subset_mapping else SubsetDataset(subset=subset_name)
+            for subset_name in subsets
+        ]
+        return [subset.set_default(dataset_meta) for subset in subsets]
+
+    def load(
+        self,
+        dataset_syntax: Optional[DatasetSyntax] = None,
+        dataset_meta: Optional[DatasetMeta] = None,
+        *,
+        use_hf: Optional[bool] = None,
+    ) -> HfDataset:
+        if dataset_syntax.dataset_type == 'path':
+            dataset = self._load_dataset_path(
+                dataset_syntax.dataset,
+                dataset_meta=dataset_meta,
+            )
+        else:
+            subsets: List[SubsetDataset] = self._select_subsets(dataset_syntax.subsets, dataset_meta)
+            revision = dataset_meta.hf_revision if use_hf else dataset_meta.ms_revision
+            datasets = []
+            for subset in subsets:
+                dataset = self._load_repo_dataset(
+                    dataset_syntax.dataset,
+                    subset,
+                    use_hf=use_hf,
+                    revision=revision,
+                )
+                datasets.append(dataset)
+            dataset = self.concat_datasets(datasets)
+        return dataset
+
+
+def init_self_cognition_preprocessor(
+    dataset_meta: Optional[DatasetMeta],
+    model_name: Optional[Union[Tuple[str, str], List[str]]] = None,
+    model_author: Optional[Union[Tuple[str, str], List[str]]] = None,
+) -> None:
+    from .dataset.llm import SelfCognitionPreprocessor
+    if dataset_meta is None or model_name is None and model_author is None:
+        return
+    kwargs = {}
+    # zh, en
+    for key in ['name', 'author']:
+        val = locals()[f'model_{key}']
+        if isinstance(val, str):
+            val = [val]
+        if val is not None and val[0] is not None and (len(val) == 1 or val[1] is None):
+            val = (val[0], val[0])
+        kwargs[key] = val
+
+    preprocess_funcs = [dataset_meta.preprocess_func]
+    preprocess_funcs += [subset.preprocess_func for subset in dataset_meta.subsets if isinstance(subset, SubsetDataset)]
+    for preprocess_func in preprocess_funcs:
+        if isinstance(preprocess_func, SelfCognitionPreprocessor):
+            preprocess_func.set_name_author(**kwargs)
+    logger.info_once(f"SelfCognitionPreprocessor has been successfully configured with name: {kwargs['name']}, "
+                     f"author: {kwargs['author']}.")
+
+
+def _inject_dataset_routing_tag(dataset: DATASET_TYPE, ds_name: str) -> DATASET_TYPE:
+    """Inject ``dataset`` column for multi-teacher routing (constant per source dataset)."""
+    if isinstance(dataset, HfIterableDataset):
+        return dataset.map(lambda example: {**example, 'dataset': ds_name})
+    return dataset.add_column('dataset', [ds_name] * len(dataset))
+
+
+def load_dataset(
+    datasets: Union[List[str], str],
+    *,
+    split_dataset_ratio: float = 0.,
+    seed: Union[int, np.random.RandomState, None] = 42,
+    num_proc: int = 1,
+    load_from_cache_file: bool = True,
+    shuffle: bool = False,
+    streaming: bool = False,
+    interleave_prob: Optional[List[float]] = None,
+    stopping_strategy: Literal['first_exhausted', 'all_exhausted'] = 'first_exhausted',
+    shuffle_buffer_size: int = 1000,
+    use_hf: Optional[bool] = None,
+    hub_token: Optional[str] = None,
+    strict: bool = False,
+    download_mode: Literal['force_redownload', 'reuse_dataset_if_exists'] = 'reuse_dataset_if_exists',
+    columns: Optional[Dict[str, str]] = None,  # columns_mapping
+    remove_unused_columns: bool = True,
+    disable_auto_column_mapping: bool = False,
+    # self-cognition
+    model_name: Optional[Union[Tuple[str, str], List[str]]] = None,  # zh, en
+    model_author: Optional[Union[Tuple[str, str], List[str]]] = None,
+) -> Tuple[DATASET_TYPE, Optional[DATASET_TYPE]]:
+    """Load and preprocess datasets.
+
+    This function provides a unified interface to load datasets from various sources (HuggingFace,
+    ModelScope, or local paths), with support for splitting, shuffling, streaming, and interleaving
+    multiple datasets. It also handles self-cognition dataset preprocessing for model training.
+
+    Args:
+        datasets: Single dataset name or list of dataset names to load. Can use special syntax
+            for advanced configurations (e.g., 'dataset_name#1000' for sampling).
+        split_dataset_ratio: Ratio for splitting dataset into train/validation sets.
+            Value between 0 and 1. If 0, no validation split is created. Default: 0.
+        seed: Random seed for reproducibility. Can be an integer or numpy RandomState object.
+            If None, results will be non-deterministic. Default: 42.
+        num_proc: Number of processes to use for dataset preprocessing. Set to None for
+            streaming mode. Default: 1.
+        load_from_cache_file: Whether to load preprocessed data from cache if available.
+            Default: True.
+        shuffle: Whether to shuffle the dataset(s) after loading. Default: False.
+        streaming: Enable streaming mode for large datasets that don't fit in memory.
+            When True, num_proc is automatically set to None. Default: False.
+        interleave_prob: Probability weights for interleaving multiple datasets. Must have
+            same length as datasets list. If None, datasets are concatenated instead. Default: None.
+        stopping_strategy: Strategy when interleaving datasets of different lengths:
+            - 'first_exhausted': Stop when shortest dataset is exhausted
+            - 'all_exhausted': Continue until all datasets are exhausted
+            Default: 'first_exhausted'.
+        shuffle_buffer_size: Buffer size for shuffling in streaming mode. Larger values
+            provide better randomization but use more memory. Default: 1000.
+        use_hf: Force using HuggingFace Hub (True) or ModelScope (False). If None,
+            it is controlled by the environment variable `USE_HF`, which defaults to '0'.
+            Default: None.
+        hub_token: Authentication token for accessing private datasets on the hub. Default: None.
+        strict: If True, raise exceptions when encountering malformed data rows.
+            If False, skip invalid rows with warnings. Default: False.
+        download_mode: How to handle existing cached datasets:
+            - 'reuse_dataset_if_exists': Use cached version if available
+            - 'force_redownload': Always download fresh copy
+            Default: 'reuse_dataset_if_exists'.
+        columns: Manual column name mapping for datasets. Dictionary mapping source column
+            names to target column names (e.g., {'text': 'content'}). Default: None.
+        remove_unused_columns: Whether to remove columns not used in preprocessing.
+            Helps reduce memory usage. Default: True.
+        disable_auto_column_mapping: By default, column names in the dataset are automatically
+            mapped. This parameter disables that behavior
+            (the `columns` parameter remains effective), defaulting to `False`.
+        model_name: Model name for self-cognition task preprocessing. Can be a tuple of
+            (Chinese_name, English_name) or list of names. Default: None.
+        model_author: Model author for self-cognition task preprocessing. Can be a tuple of
+            (Chinese_author, English_author) or list of authors. Default: None.
+
+    Returns:
+        A tuple of (train_dataset, val_dataset):
+            - train_dataset: The training dataset
+            - val_dataset: The validation dataset if split_dataset_ratio > 0, otherwise None
+
+    Examples:
+        >>> # Load single dataset
+        >>> train_ds, val_ds = load_dataset('AI-ModelScope/alpaca-gpt4-data-zh', split_dataset_ratio=0.1)
+
+        >>> # Load multiple datasets
+        >>> train_ds, _ = load_dataset(
+        ...     ['AI-ModelScope/alpaca-gpt4-data-zh#500', 'swift/self-cognition#500'],
+        ...     model_name=('我的模型', 'MyModel'),
+        ...     model_author=('作者', 'Author')
+        ... )
+    """
+    init_self_cognition_preprocessor(DATASET_MAPPING.get('self-cognition'), model_name, model_author)
+    if isinstance(datasets, str):
+        datasets = [datasets]
+    if not isinstance(seed, np.random.RandomState):
+        seed = np.random.RandomState(seed)
+    if streaming:
+        num_proc = None
+    train_datasets = []
+    val_datasets = []
+
+    use_hf_default = use_hf
+    if use_hf_default is None:
+        use_hf_default = True if use_hf_hub() else False
+    for dataset in datasets:
+        dataset_syntax = DatasetSyntax.parse(dataset)
+        use_hf = dataset_syntax.use_hf or use_hf_default
+        # compat dataset_name
+        if dataset_syntax.dataset in DATASET_MAPPING:
+            dataset_meta = DATASET_MAPPING[dataset_syntax.dataset]
+            if dataset_syntax.use_hf is None and dataset_meta.dataset_path is not None:
+                dataset_syntax.dataset = dataset_meta.dataset_path
+                dataset_syntax.dataset_type = 'path'
+            else:
+                dataset_syntax.dataset = dataset_meta.hf_dataset_id if use_hf else dataset_meta.ms_dataset_id
+        else:
+            dataset_meta = dataset_syntax.get_dataset_meta(use_hf)
+        loader = dataset_meta.loader(
+            num_proc=num_proc,
+            load_from_cache_file=load_from_cache_file,
+            streaming=streaming,
+            hub_token=hub_token,
+            strict=strict,
+            download_mode=download_mode,
+            columns=columns,  # columns_mapping
+            remove_unused_columns=remove_unused_columns,
+            disable_auto_column_mapping=disable_auto_column_mapping,
+        )
+        train_dataset = loader.load(dataset_syntax, dataset_meta, use_hf=use_hf)
+        train_dataset, val_dataset = loader.post_process(
+            train_dataset,
+            dataset_sample=dataset_syntax.dataset_sample,
+            split_dataset_ratio=split_dataset_ratio,
+            streaming=streaming,
+            shuffle=shuffle,
+            random_state=seed,
+        )
+        if train_dataset is not None:
+            # Inject dataset_syntax.dataset as routing tag for multi-teacher
+            ds_name = dataset_syntax.dataset
+            train_dataset = _inject_dataset_routing_tag(train_dataset, ds_name)
+            train_datasets.append(train_dataset)
+        if val_dataset is not None:
+            ds_name = dataset_syntax.dataset
+            val_dataset = _inject_dataset_routing_tag(val_dataset, ds_name)
+            val_datasets.append(val_dataset)
+
+    if interleave_prob is None:
+        train_datasets = loader.concat_datasets(train_datasets)
+        val_datasets = loader.concat_datasets(val_datasets)
+    else:
+        train_datasets = loader.interleave_datasets(
+            train_datasets, interleave_prob, seed=get_seed(seed), stopping_strategy=stopping_strategy)
+        val_datasets = loader.interleave_datasets(
+            val_datasets, interleave_prob, seed=get_seed(seed), stopping_strategy=stopping_strategy)
+
+    if shuffle:
+        if train_datasets:
+            train_datasets = loader.shuffle_dataset(
+                train_datasets, seed=get_seed(seed), buffer_size=shuffle_buffer_size)
+        if val_datasets:
+            val_datasets = loader.shuffle_dataset(val_datasets, seed=get_seed(seed), buffer_size=shuffle_buffer_size)
+    return train_datasets, val_datasets

swift/dataset/media.py

@@ -0,0 +1,128 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import aiohttp
+import os
+import shutil
+from modelscope.hub.utils.utils import get_cache_dir
+from typing import List, Literal, Optional, Union
+
+from swift.utils import get_logger, safe_ddp_context
+
+logger = get_logger()
+
+
+class MediaResource:
+    """A class to manage the resource downloading."""
+
+    cache_dir = os.path.join(get_cache_dir(), 'media_resources')
+    lock_dir = os.path.join(get_cache_dir(), 'lockers')
+
+    media_type_urls = {
+        'llava', 'coco', 'sam', 'gqa', 'ocr_vqa', 'textvqa', 'VG_100K', 'VG_100K_2', 'share_textvqa', 'web-celebrity',
+        'web-landmark', 'wikiart'
+    }
+
+    URL_PREFIX = 'https://www.modelscope.cn/api/v1/datasets/hjh0119/sharegpt4v-images/repo?Revision=master&FilePath='
+
+    @staticmethod
+    def get_url(media_type):
+        is_ocr_vqa = (media_type == 'ocr_vqa')
+        extension = 'tar' if is_ocr_vqa else 'zip'
+        return f'{MediaResource.URL_PREFIX}{media_type}.{extension}'
+
+    @staticmethod
+    def download(media_type_or_url: Union[str, List[str]],
+                 local_alias: Optional[str] = None,
+                 file_type: Literal['compressed', 'file', 'sharded'] = 'compressed'):
+        """Download and extract a resource from a http link.
+
+        Args:
+            media_type_or_url: `str` or List or `str`, Either belongs to the `media_type_urls`
+                listed in the class field, or a remote url to download and extract.
+                Be aware that, this media type or url needs to contain a zip or tar file.
+            local_alias: `Options[str]`, The local alias name for the `media_type_or_url`. If the first arg is a
+                media_type listed in this class, local_alias can leave None. else please pass in a name for the url.
+                The local dir contains the extracted files will be: {cache_dir}/{local_alias}
+            file_type: The file type, if is a compressed file, un-compressed the file,
+                if is an original file, only download it, if is a sharded file, download all files and extract.
+
+        Returns:
+            The local dir contains the extracted files.
+        """
+        media_file = media_type_or_url if isinstance(media_type_or_url, str) else media_type_or_url[0]
+        with safe_ddp_context(hash_id=media_file):
+            return MediaResource._safe_download(
+                media_type=media_type_or_url, media_name=local_alias, file_type=file_type)
+
+    @staticmethod
+    def move_directory_contents(src_dir, dst_dir):
+        if not os.path.exists(dst_dir):
+            os.makedirs(dst_dir)
+
+        for dirpath, dirnames, filenames in os.walk(src_dir):
+            relative_path = os.path.relpath(dirpath, src_dir)
+            target_dir = os.path.join(dst_dir, relative_path)
+
+            if not os.path.exists(target_dir):
+                os.makedirs(target_dir)
+
+            for file in filenames:
+                src_file = os.path.join(dirpath, file)
+                dst_file = os.path.join(target_dir, file)
+                shutil.move(src_file, dst_file)
+
+    @staticmethod
+    def _safe_download(media_type: Union[str, List[str]],
+                       media_name: Optional[str] = None,
+                       file_type: Literal['compressed', 'file', 'sharded'] = 'compressed'):
+        media_name = media_name or media_type
+        assert isinstance(media_name, str), f'{media_name} is not a str'
+        if isinstance(media_type, str) and media_type in MediaResource.media_type_urls:
+            media_type = MediaResource.get_url(media_type)
+
+        from datasets.download.download_manager import DownloadConfig, DownloadManager
+        final_folder = os.path.join(MediaResource.cache_dir, media_name)
+
+        if file_type == 'file':
+            filename = media_type.split('/')[-1]
+            final_path = os.path.join(final_folder, filename)
+            if os.path.exists(final_path):  # if the download thing is a file but not folder,
+                return final_folder  # check whether the file exists
+            if not os.path.exists(final_folder):
+                os.makedirs(final_folder)  # and make sure final_folder exists to contain it
+        else:
+            if os.path.exists(final_folder):
+                return final_folder
+
+        logger.info('# #################Resource downloading#################')
+        logger.info('Downloading necessary resources...')
+        logger.info(f'Resource package: {media_type}')
+        logger.info(f'Extracting to local dir: {final_folder}')
+        logger.info('If the downloading fails or lasts a long time, '
+                    'you can manually download the resources and extracting to the local dir.')
+        logger.info('Now begin.')
+        download_config = DownloadConfig(cache_dir=MediaResource.cache_dir)
+        download_config.storage_options = {'client_kwargs': {'timeout': aiohttp.ClientTimeout(total=86400)}}
+        if file_type == 'file':
+            filename = media_type.split('/')[-1]
+            final_path = os.path.join(final_folder, filename)
+            local_dirs = DownloadManager(download_config=download_config).download(media_type)
+            shutil.move(str(local_dirs), final_path)
+        elif file_type == 'compressed':
+            local_dirs = DownloadManager(download_config=download_config).download_and_extract(media_type)
+            shutil.move(str(local_dirs), final_folder)
+        else:
+            for media_url in media_type:
+                local_dirs = DownloadManager(download_config=download_config).download_and_extract(media_url)
+                MediaResource.move_directory_contents(str(local_dirs), final_folder)
+        logger.info('# #################Resource downloading finished#################')
+        return final_folder
+
+    @staticmethod
+    def safe_save(image, file_name, folder, format='JPEG'):
+        folder = os.path.join(MediaResource.cache_dir, folder)
+        os.makedirs(folder, exist_ok=True)
+        file = os.path.join(folder, file_name)
+        if os.path.exists(file):
+            return file
+        image.save(file, format=format)
+        return file

swift/dataset/packing.py

@@ -0,0 +1,230 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import math
+import multiprocessing as mp
+import torch.distributed as dist
+from itertools import chain
+from torch.utils.data import Dataset, IterableDataset
+from tqdm import tqdm
+from typing import Optional
+
+from swift.template import MaxLengthError
+from swift.utils import get_logger, is_dist, is_master, split_list
+
+logger = get_logger()
+
+
+def calculate_matched_group(sequences, packing_length: int, is_finished: bool = True, strategy: str = 'binpack'):
+    if len(sequences) == 0:
+        return [], []
+    if strategy == 'sequential':
+        # Order-preserving greedy packing (next-fit): keep a single open pack and flush it
+        # when the next sample doesn't fit, so the global sample order and pack boundaries
+        # follow the input order (a sequential sampler). (Use packing_num_proc=1 for
+        # a single global ordering.)
+        packs, cur, cur_len = [], [], 0
+        for item in sequences:  # item = (idx, length); weight_pos=1 -> length at item[1]
+            seq_len = item[1]
+            if cur and cur_len + seq_len > packing_length:
+                packs.append(cur)
+                cur, cur_len = [], 0
+            cur.append(item)
+            cur_len += seq_len
+            if cur_len >= packing_length:
+                packs.append(cur)
+                cur, cur_len = [], 0
+        if is_finished:
+            if cur:
+                packs.append(cur)
+            return packs, []
+        return packs, cur
+    # default: best-fit-decreasing bin packing (https://arxiv.org/pdf/2404.10830)
+    import binpacking
+    sequences = binpacking.to_constant_volume(sequences, packing_length, weight_pos=1)
+    if sequences and not is_finished:
+        sequences, ret_sequences = sequences[:-1], sequences[-1]
+    else:
+        ret_sequences = []
+    return sequences, ret_sequences
+
+
+class PackingDataset(Dataset):
+    PACKING_BATCH_SIZE = 1000
+
+    def __init__(
+        self,
+        template,
+        dataset,
+        num_proc: int = 1,
+        *,
+        strict: bool = False,
+        load_from_cache_file: bool = True,
+        packing_length: Optional[int] = None,
+        packing_num_proc: int = 1,
+        packing_strategy: str = 'binpack',
+        **kwargs,
+    ):
+        template.packing = True
+        template.padding_free = True  # TODO: remove
+        self.template = template
+        self.dataset = dataset
+        self.num_proc = num_proc
+        self.strict = strict
+        self.load_from_cache_file = load_from_cache_file
+        self.packing_strategy = packing_strategy
+        self.packing_length = packing_length or self.template.max_length
+        self.packing_num_proc = min(packing_num_proc, math.ceil(len(dataset) / self.PACKING_BATCH_SIZE))
+        self._out_queue = mp.Queue()
+        if is_master():
+            lengths = self.dataset['lengths']
+            offset = 0
+            chunked_lengths = split_list(lengths, self.packing_num_proc)
+            for i in range(self.packing_num_proc):
+                worker = mp.Process(
+                    target=self.create_packed_idx, args=(
+                        i,
+                        offset,
+                        chunked_lengths[i],
+                    ), daemon=True)
+                worker.start()
+                offset += len(chunked_lengths[i])
+            self.packed_idx = [[] for _ in range(self.packing_num_proc)]
+            self.packed_length = [[] for _ in range(self.packing_num_proc)]
+            desc = 'Packing: ' if self.packing_num_proc == 1 else f'Packing (num_proc={self.packing_num_proc}): '
+            with tqdm(total=len(lengths), dynamic_ncols=True, desc=desc) as prog_bar:
+                finished_workers = 0
+                while finished_workers < self.packing_num_proc:
+                    rank, sequences, data_len = self._out_queue.get()
+                    if data_len == -1:
+                        finished_workers += 1
+                        continue
+                    prog_bar.update(data_len)
+                    self.packed_idx[rank] += [[x[0] for x in seq] for seq in sequences]
+                    self.packed_length[rank] += [sum(x[1] for x in seq) for seq in sequences]
+            self.packed_idx = list(chain.from_iterable(self.packed_idx))
+            self.packed_length = list(chain.from_iterable(self.packed_length))
+        else:
+            self.packed_idx, self.packed_length = None, None
+        if dist.is_initialized() and is_dist():
+            obj_list = [(self.packed_idx, self.packed_length)]
+            dist.broadcast_object_list(obj_list)
+            self.packed_idx, self.packed_length = obj_list[0]
+
+    def create_packed_idx(self, rank, offset, lengths):
+        data = [(i + offset, sum(length) if isinstance(length, list) else length) for i, length in enumerate(lengths)]
+        i = 0
+        input_data = []
+        while True:
+            new_data = data[i:i + self.PACKING_BATCH_SIZE]
+            input_data += new_data
+            if not input_data:
+                break
+            i += self.PACKING_BATCH_SIZE
+            is_finished = i >= len(data)
+            sequences, input_data = calculate_matched_group(
+                input_data, self.packing_length, is_finished=is_finished, strategy=self.packing_strategy)
+            self._out_queue.put((rank, sequences, len(new_data)))
+        self._out_queue.put((rank, [], -1))
+
+    def __getitem__(self, index):
+        sequence = self.packed_idx[index]
+        row = [self.dataset[i] for i in sequence]
+        return row
+
+    def __len__(self):
+        return len(self.packed_idx)
+
+
+class IterablePackingDataset(IterableDataset):
+
+    def __init__(
+        self,
+        template,
+        dataset,
+        num_proc: int = 1,
+        *,
+        packing_interval: int = 128,
+        packing_length: Optional[int] = None,
+        strict: bool = False,
+        cyclic: bool = False,
+        packing_strategy: str = 'binpack',
+        **kwargs,
+    ):
+        template.packing = True
+        template.padding_free = True  # TODO: remove
+        self.template = template
+        self.dataset = dataset
+        self.num_proc = num_proc
+        self.strict = strict
+        self.packing_length = packing_length or self.template.max_length
+
+        self.packing_interval = packing_interval
+        self._in_queue = mp.Queue()
+        self._out_queue = mp.Queue()
+        self.workers = []
+        self.cyclic = cyclic
+        self.packing_strategy = packing_strategy
+        for _ in range(self.num_proc):
+            worker = mp.Process(target=self._processor, daemon=True)
+            worker.start()
+            self.workers.append(worker)
+
+    def _processor(self):
+        while True:
+            i, data = self._in_queue.get()
+            encoded_data = {}
+            try:
+                encoded_data = self.template.encode(data, return_length=True)
+            except Exception as e:
+                if self.strict and not isinstance(e, MaxLengthError):
+                    raise
+            self._out_queue.put((i, encoded_data))
+
+    def _put_data_in_queue(self, iterator) -> int:
+        for i in range(self.packing_interval):
+            try:
+                data = next(iterator)
+            except StopIteration:
+                return i
+            self._in_queue.put((i, data))
+        return i + 1
+
+    def _fetch_data_out_queue(self, last_res, num_samples):
+        res = [None] * num_samples
+        for _ in range(num_samples):
+            i, data = self._out_queue.get()
+            if not data:
+                continue
+            res[i] = (data, len(data['input_ids']))
+        res = [data for data in res if data]
+        last_res += res
+        return last_res
+
+    @staticmethod
+    def cyclic_iter(iterable):
+        while True:
+            for x in iterable:
+                yield x
+
+    def __iter__(self):
+        try:
+            next(iter(self.dataset))
+        except StopIteration:
+            return
+
+        if self.cyclic:
+            iterator = self.cyclic_iter(self.dataset)
+        else:
+            iterator = iter(self.dataset)
+        data = []
+        while True:
+            num_samples = self._put_data_in_queue(iterator)
+            finished = num_samples != self.packing_interval
+            data = self._fetch_data_out_queue(data, num_samples)
+            sequences, data = calculate_matched_group(
+                data, self.packing_length, is_finished=finished, strategy=self.packing_strategy)
+            res = []
+            for row in sequences:
+                res.append([r[0] for r in row])
+            yield from res
+            if finished:
+                break

swift/dataset/preprocessor/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from .core import (DATASET_TYPE, AlpacaPreprocessor, AutoPreprocessor, ClsPreprocessor, MessagesPreprocessor,
+                   ResponsePreprocessor, RowPreprocessor)
+from .extra import ClsGenerationPreprocessor, GroundingMixin, TextGenerationPreprocessor

swift/dataset/preprocessor/core.py

@@ -0,0 +1,571 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import ast
+import datasets
+import numpy as np
+import os
+from collections import Counter
+from contextlib import contextmanager
+from datasets import Dataset as HfDataset
+from datasets import Image
+from datasets import IterableDataset as HfIterableDataset
+from datasets import Sequence, Value
+from itertools import chain
+from modelscope.hub.utils.utils import get_cache_dir
+from packaging import version
+from typing import Any, Callable, Dict, List, Optional, Union
+
+from swift.template import history_to_messages
+from swift.utils import get_logger, is_dist, is_master, safe_ddp_context
+
+DATASET_TYPE = Union[HfDataset, HfIterableDataset]
+
+logger = get_logger()
+
+_pair_keys = ['messages', 'images', 'videos', 'audios', 'tools', 'objects']
+
+
+class RowPreprocessor:
+    standard_keys = _pair_keys + list(
+        chain.from_iterable([f'{prefix}_{k}' for k in _pair_keys]
+                            for prefix in ['rejected', 'positive', 'negative'])) + [
+                                'rejected_response',
+                                'label',
+                                'channel',
+                                'margin',
+                                'teacher_prompt',
+                                'chat_template_kwargs',
+                                # Qwen3-TTS
+                                'ref_audios',
+                                'audio_codes',
+                            ]
+
+    def __init__(self,
+                 *,
+                 columns: Optional[Dict[str, str]] = None,
+                 dataset_sample: Optional[int] = None,
+                 random_state: Optional[Union[np.random.RandomState, int]] = 42,
+                 traceback_limit: int = 10) -> None:
+        self.columns = columns or {}
+        self.origin_columns = self.columns.copy()  # Higher priority and raise Error
+        images_keys = ['images', 'image']
+        audios_keys = ['audios', 'audio']
+        videos_keys = ['videos', 'video']
+        for mm_type in ['images', 'audios', 'videos']:
+            keys = locals()[f'{mm_type}_keys']
+            for key in keys:
+                self.columns[key] = mm_type
+
+        self.traceback_limit = traceback_limit
+        self._traceback_counter = 0
+        self.dataset_sample = dataset_sample
+        self.datasets_4 = version.parse(datasets.__version__) >= version.parse('4.0')
+        if not isinstance(random_state, np.random.RandomState):
+            random_state = np.random.RandomState(random_state)
+        self.random_state = random_state
+
+    @staticmethod
+    def _check_messages(row: Dict[str, Any]) -> None:
+        if 'messages' not in row:
+            return
+        messages = row['messages']
+        assert len(messages) > 0, f'messages: {messages}'
+        # fix swift/SlimOrca (concat)
+        for message in messages:
+            keys = set(message.keys()) - {'role', 'content', 'loss', 'loss_scale'}
+            for key in keys:
+                message.pop(key)
+
+        for message in messages:
+            role, content = message['role'], message['content']
+            # The terms 'tool' and 'tool_response' have the same meaning, ensuring compatibility.
+            assert role in {'system', 'user', 'tool_call', 'tool_response', 'tool', 'assistant'}, f'message: {message}'
+            assert content is not None, f'message: {message}'
+
+    @staticmethod
+    def _cast_mm_data(row: Dict[str, Any]) -> None:
+        for key in ['images', 'rejected_images']:
+            images = row.get(key, None)
+            if images is None:
+                continue
+
+            if isinstance(images, str) or (isinstance(images, list) and images and isinstance(images[0], str)):
+                if isinstance(images, str):
+                    images = [images]
+                for i, image in enumerate(images):
+                    images[i] = {'bytes': None, 'path': image}
+                row[key] = images
+            elif isinstance(images, dict):
+                row[key] = [images]
+
+        for key in ['videos', 'audios']:
+            mm_data = row.get(key)
+            if mm_data is None:
+                continue
+            elif isinstance(mm_data, str):
+                row[key] = [mm_data]
+
+    @staticmethod
+    def _check_rejected_response(row: Dict[str, Any]) -> None:
+        if 'rejected_response' in row:
+            messages = row['messages']
+            rejected_response = row['rejected_response']
+            if (rejected_response is None
+                    or isinstance(rejected_response, str) and rejected_response == messages[-1]['content']):
+                raise ValueError(f'rejected_response: {rejected_response}')
+
+    def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        return row
+
+    def prepare_dataset(self, dataset: DATASET_TYPE) -> DATASET_TYPE:
+        return dataset
+
+    @staticmethod
+    def batched_to_rows(batched_row: Dict[str, Any]):
+        keys = list(batched_row.keys())
+        batch_size = len(batched_row[keys[0]])
+        return [{key: batched_row[key][i] for key in keys} for i in range(batch_size)]
+
+    @staticmethod
+    def rows_to_batched(rows: List[Dict[str, Any]]):
+        batched = {}
+        for i, row in enumerate(rows):
+            for k, v in row.items():
+                if k not in batched:
+                    batched[k] = [None] * i
+                batched[k].append(v)
+            # Make all the lengths of v the same.
+            for k in set(batched.keys()) - set(row.keys()):
+                batched[k].append(None)
+        return batched
+
+    @staticmethod
+    def _remove_prefix_keys(row, prefix: str):
+        for k in list(row.keys()):
+            if k.startswith(prefix):
+                new_k = k[len(prefix):]
+                new_v = row.pop(k)
+                if new_k not in row:
+                    row[new_k] = new_v
+
+    @staticmethod
+    def _check_objects(row):
+        objects = row.get('objects')
+        if objects is None:
+            return
+        new_objects = {}
+        # Ensure the order
+        for k in ['ref', 'bbox', 'bbox_type', 'image_id']:
+            if k in objects.keys():
+                new_objects[k] = objects[k]
+        row['objects'] = new_objects
+        bbox = new_objects['bbox']
+
+        # check bbox
+        for box in bbox:
+            assert len(box) in {2, 4}, f'len(box): {len(box)}'
+            if len(box) == 2:
+                continue
+            if box[0] > box[2]:
+                box[0], box[2] = box[2], box[0]
+            if box[1] > box[3]:
+                box[1], box[3] = box[3], box[1]
+
+    def batched_preprocess(self, batched_row: Dict[str, Any], *, strict: bool,
+                           ignore_max_length_error: bool) -> Dict[str, Any]:
+        from swift.template import MaxLengthError
+        batched_row = dict(batched_row)
+        assert len(batched_row) > 0
+        self._remove_prefix_keys(batched_row, '__@')  # compat streaming
+        rows = self.batched_to_rows(batched_row)
+
+        new_rows = []
+        for row in rows:
+            try:
+                row = self.preprocess(row)
+                # support [row1, row2, ...]
+                if row is None:
+                    row = []
+                if isinstance(row, dict):
+                    row = [row]
+                for r in row:
+                    self._check_objects(r)
+                    self._check_rejected_response(r)
+                    self._check_messages(r)
+                    self._cast_mm_data(r)
+            except Exception as e:
+                if strict:
+                    logger.warning('To avoid errors, you can pass `strict=False`.')
+                    raise
+                if isinstance(e, MaxLengthError) and ignore_max_length_error:
+                    pass
+                elif self.traceback_limit is not None and self._traceback_counter < self.traceback_limit:
+                    import traceback
+                    logger.info(traceback.format_exc())
+                    logger.warning('👆👆👆There are errors in the dataset, the data will be deleted')
+                    self._traceback_counter += 1
+                row = []
+            new_rows += row
+        res = self.rows_to_batched(new_rows)
+        self._remove_prefix_keys(res, '__#')  # compat GRPO
+        if len(res) == 0:
+            res['messages'] = []
+
+        return res
+
+    @staticmethod
+    def get_features_dataset(dataset: DATASET_TYPE) -> DATASET_TYPE:
+        if dataset.features is None:
+            assert isinstance(dataset, HfIterableDataset)
+            dataset = dataset._resolve_features()
+        return dataset
+
+    @staticmethod
+    def safe_rename_columns(dataset, columns):
+        dataset = RowPreprocessor.get_features_dataset(dataset)
+        columns_keys = {k.lower(): k for k in dataset.features.keys()}  # lower -> lower/upper
+        safe_columns = {columns_keys[k.lower()]: v for k, v in columns.items() if k.lower() in columns_keys}
+
+        counter = Counter(safe_columns.values())
+        for k, new_k in list(safe_columns.items()):
+            if counter[new_k] > 1:
+                # For example, if "response" and "answer" match, then no processing is done.
+                safe_columns.pop(k)
+                continue
+
+        # e.g. Keep {'query': 'query'} to ensure that the query has the highest priority.
+        safe_columns = {k: v for k, v in safe_columns.items() if k != v}
+        if safe_columns:
+            dataset = dataset.rename_columns(safe_columns)
+
+        return dataset
+
+    @staticmethod
+    def remove_useless_columns(dataset: DATASET_TYPE) -> DATASET_TYPE:
+        dataset = RowPreprocessor.get_features_dataset(dataset)
+        features = dataset.features
+        k_list = [k for k in RowPreprocessor.standard_keys if k in features]
+        if len(k_list) != len(features):
+            dataset = dataset.select_columns(k_list)
+        return dataset
+
+    @contextmanager
+    def _patch_arrow_writer(self):
+        # fix AI-ModelScope/ms_agent_for_agentfabric:all
+        from datasets.arrow_writer import ArrowWriter
+
+        def _new_init(_self, schema=None, features=None, *args, **kwargs):
+
+            if features is not None:
+
+                if self.datasets_4:
+                    from datasets.features import Json, List
+                    messages_feature = List(Json())
+                    for key in ['messages', 'rejected_messages', 'positive_messages', 'negative_messages']:
+                        features[key] = messages_feature
+                    features['images'] = List({'bytes': Value(dtype='binary'), 'path': Value(dtype='string')})
+                    features['objects'] = Json()
+                    features['chat_template_kwargs'] = Json()
+                else:
+                    messages_feature = [{
+                        'role': Value(dtype='string'),
+                        'content': Value(dtype='string'),
+                    }]
+                    messages_feature_with_loss = [{
+                        'role': Value(dtype='string'),
+                        'content': Value(dtype='string'),
+                        'loss': Value(dtype='bool'),
+                        'loss_scale': Value(dtype='float64'),
+                    }]
+                    features['messages'] = messages_feature_with_loss
+                    features['rejected_messages'] = messages_feature_with_loss
+                    features['positive_messages'] = messages_feature
+                    features['negative_messages'] = messages_feature
+                    features['images'] = [{'bytes': Value(dtype='binary'), 'path': Value(dtype='string')}]
+                    features['objects'] = {
+                        'ref': Sequence(feature=Value(dtype='string'), length=-1),
+                        'bbox': Sequence(feature=Sequence(feature=Value(dtype='float64'), length=-1), length=-1),
+                        'bbox_type': Value(dtype='string'),
+                        'image_id': Sequence(feature=Value(dtype='int64'), length=-1),
+                    }
+            ArrowWriter.__origin_init__(_self, schema, features, *args, **kwargs)
+
+        ArrowWriter.__origin_init__ = ArrowWriter.__init__
+        ArrowWriter.__init__ = _new_init
+        try:
+            yield
+        finally:
+            ArrowWriter.__init__ = ArrowWriter.__origin_init__
+            del ArrowWriter.__origin_init__
+
+    def _cast_pil_image(self, dataset):
+        features = dataset.features
+        for col in ['images', 'rejected_images']:
+            if (col in features and isinstance(features[col], Image) and getattr(features[col], 'decode', False)):
+                dataset = dataset.cast_column(col, Image(decode=False))
+        return dataset
+
+    def __call__(
+        self,
+        dataset: DATASET_TYPE,
+        *,
+        num_proc: int = 1,
+        load_from_cache_file: bool = True,
+        strict: bool = False,
+        batch_size: Optional[int] = None,
+        enable_auto_mapping: bool = False,
+    ) -> DATASET_TYPE:
+        from ..utils import sample_dataset
+        if batch_size is None:
+            batch_size = 1000 if isinstance(dataset, HfDataset) else 16
+        if self.dataset_sample is not None:
+            dataset = sample_dataset(dataset, self.dataset_sample, True, self.random_state)
+
+        map_kwargs = {'batched': True, 'batch_size': batch_size}
+        if isinstance(dataset, HfDataset):
+            if not load_from_cache_file and is_dist() and not is_master():
+                load_from_cache_file = True
+            map_kwargs.update({
+                'num_proc': num_proc,
+                'load_from_cache_file': load_from_cache_file,
+            })
+        # compat GRPO: The solution field will be retained.
+        dataset = RowPreprocessor.get_features_dataset(dataset)
+        if 'solution' in dataset.features:
+            with safe_ddp_context(None, True):
+                if isinstance(dataset, HfDataset) and not dataset.cache_files:
+                    map_kwargs['cache_file_name'] = os.path.join(get_cache_dir(), 'datasets', 'map_cache',
+                                                                 f'{dataset._fingerprint}.arrow')
+                dataset = dataset.map(lambda x: {'__#solution': x['solution']}, **map_kwargs)
+                map_kwargs.pop('cache_file_name', None)
+        dataset = self.safe_rename_columns(dataset, self.origin_columns)
+        if enable_auto_mapping:
+            dataset = self.safe_rename_columns(dataset, self.columns)
+        dataset = self.prepare_dataset(dataset)
+        dataset = self._cast_pil_image(dataset)
+        if isinstance(dataset, HfIterableDataset):
+            # fix: https://github.com/huggingface/datasets/issues/6408
+            columns = {k: f'__@{k}' for k in RowPreprocessor.standard_keys if k in dataset.features}
+            if columns:
+                dataset = dataset.rename_columns(columns)
+
+        ignore_max_length_error = True
+        with self._patch_arrow_writer(), safe_ddp_context(None, True):
+            if isinstance(dataset, HfDataset) and not dataset.cache_files:
+                map_kwargs['cache_file_name'] = os.path.join(get_cache_dir(), 'datasets', 'map_cache',
+                                                             f'{dataset._fingerprint}.arrow')
+            dataset_mapped = dataset.map(
+                self.batched_preprocess,
+                fn_kwargs={
+                    'strict': strict,
+                    'ignore_max_length_error': ignore_max_length_error,
+                },
+                remove_columns=list(dataset.features.keys()),
+                **map_kwargs)
+        if isinstance(dataset_mapped, HfDataset) and len(dataset) != len(dataset_mapped):
+            logger.info(
+                f'Dataset filtered, origin length: {len(dataset)}, filtered dataset length: {len(dataset_mapped)}')
+
+        return dataset_mapped
+
+
+class ResponsePreprocessor(RowPreprocessor):
+    """Dataset compatible with older versions of ms-swift"""
+
+    def __init__(self, *, columns: Optional[Dict[str, str]] = None, **kwargs) -> None:
+        super().__init__(columns=columns, **kwargs)
+        system_keys = ['system', 'system_prompt']
+        query_keys = ['query', 'prompt', 'input', 'instruction', 'question', 'problem']
+        response_keys = ['response', 'answer', 'output', 'targets', 'target', 'answer_key', 'answers', 'solution'
+                         ] + ['text', 'completion', 'content']
+        for key in system_keys:
+            self.columns[key] = 'system'
+        for key in query_keys:
+            self.columns[key] = 'query'
+        for key in response_keys:
+            self.columns[key] = 'response'
+
+    def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        response = row.pop('response', None)
+        if response is not None:
+            if isinstance(response, (list, tuple)):
+                from transformers.utils import strtobool
+
+                # sometimes response is a list, pick one randomly
+                if strtobool(os.environ.get('RANDOM_DATASET_RESPONSE', 'False')):
+                    response = self.random_state.choice(response)
+                else:
+                    response = response[0]
+        history = row.pop('history', None) or []
+        query = row.pop('query', None)
+        system = row.pop('system', None)
+        if isinstance(history, str):  # e.g. "[['query1', 'response1']]"
+            history = ast.literal_eval(history)
+        history.append([query, response])
+
+        row.update({'messages': history_to_messages(history, system)})
+        return row
+
+
+class AlpacaPreprocessor(ResponsePreprocessor):
+
+    @classmethod
+    def concat_inst_input(cls, instruction, input_):
+        if instruction and input_:
+            query = f'{instruction}\n{input_}'
+        else:
+            query = instruction or input_
+        assert isinstance(query, str), f'query: {query}'
+        return query
+
+    def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        instruction = row.pop('instruction', None)
+        input_ = row.pop('input', None)
+        output = row.pop('output', None)
+        if output is not None:
+            row['response'] = output
+        row['query'] = self.concat_inst_input(instruction, input_)
+        return super().preprocess(row)
+
+
+def default_repair_messages(s: Union[str, Any]) -> Any:
+    if isinstance(s, str):
+        return ast.literal_eval(s)
+    return s
+
+
+class MessagesPreprocessor(RowPreprocessor):
+
+    def __init__(
+            self,
+            *,
+            # If set to None, automatic matching will be performed.
+            role_key: Optional[str] = None,  # 'role', 'from'
+            content_key: Optional[str] = None,  # 'content', 'value'
+            user_role: Optional[str] = None,  # 'user', 'human'
+            assistant_role: Optional[str] = None,  # 'assistant', 'gpt', 'bot'
+            system_role: str = 'system',
+            # 'conversation', 'conversations' -> 'messages'
+            columns: Optional[Dict[str, str]] = None,
+            repair_messages: Callable[[Union[str, List[Dict[str, str]]]],
+                                      Optional[List[Dict[str, str]]]] = default_repair_messages,
+            inner_key: Optional[str] = None,
+            **kwargs):
+        super().__init__(columns=columns, **kwargs)
+        self.role_keys = ['role', 'from'] if role_key is None else [role_key]
+        self.content_keys = ['content', 'value'] if content_key is None else [content_key]
+        self.user_roles = ['user', 'human'] if user_role is None else [user_role]
+        self.assistant_roles = ['assistant', 'gpt', 'bot'] if assistant_role is None else [assistant_role]
+        self.tool_call_roles = ['function_call']
+        self.tool_response_roles = ['function_response', 'observation', 'observations']
+
+        self.system_role = system_role
+        self.repair_messages = repair_messages
+        self.inner_key = inner_key
+
+        message_keys = ['messages', 'conversation', 'conversations']
+        for key in message_keys:
+            self.columns[key] = 'messages'
+        # sharegptq
+        system_keys = ['system', 'system_prompt']
+        if system_role not in system_keys:
+            system_keys.append(system_role)
+        for key in system_keys:
+            self.columns[key] = 'system'
+
+    @staticmethod
+    def _is_sharegpt_format(message: Dict[str, str]) -> bool:
+        if 'role' in message or 'content' in message:
+            return False
+        return True
+
+    def sharegpt_to_messages(self, messages: List[Dict[str, str]], system: Optional[str]) -> List[Dict[str, str]]:
+        self._to_std_key(messages, 'user', self.user_roles)
+        self._to_std_key(messages, 'assistant', self.assistant_roles)
+        new_messages = []
+        if system is not None:
+            new_messages.append({'role': 'system', 'content': system})
+        for message in messages:
+            user_message = {'role': 'user', 'content': message['user']}
+            assistant_message = {'role': 'assistant', 'content': message['assistant']}
+            new_messages.append(user_message)
+            new_messages.append(assistant_message)
+        return new_messages
+
+    def to_std_messages(self, messages: List[Dict[str, str]], system: Optional[str]) -> None:
+        if messages[0]['role'] == self.system_role:
+            messages[0]['role'] = 'system'
+        elif system is not None:
+            messages.insert(0, {'role': 'system', 'content': system})
+        for message in messages:
+            role = message['role']
+            if role in self.user_roles:
+                message['role'] = 'user'
+            elif role in self.assistant_roles:
+                message['role'] = 'assistant'
+            elif role.replace('-', '_') in self.tool_call_roles:
+                message['role'] = 'tool_call'
+            elif role.replace('-', '_') in self.tool_response_roles:
+                message['role'] = 'tool_response'
+
+    @staticmethod
+    def _to_std_key(messages: List[Dict[str, str]], std_key: str, optional_keys: List[str]) -> None:
+        for message in messages:
+            for key in optional_keys:
+                if key in message:
+                    message[std_key] = message.pop(key)
+
+    def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        if 'rejected_messages' in row:
+            rejected = MessagesPreprocessor.preprocess(self, {'messages': row['rejected_messages']})
+            row['rejected_messages'] = rejected['messages'] if rejected else None
+        messages = row['messages']
+        if self.inner_key is not None:
+            messages = messages[self.inner_key]
+        messages: Optional[List[Dict[str, str]]] = self.repair_messages(messages)
+        if not messages or isinstance(messages, str):
+            return
+        self._to_std_key(messages, 'role', self.role_keys)
+        self._to_std_key(messages, 'content', self.content_keys)
+        system = row.pop('system', None)
+        if self._is_sharegpt_format(messages[0]):
+            messages = self.sharegpt_to_messages(messages, system)
+        else:
+            self.to_std_messages(messages, system)  # inplace
+        row['messages'] = messages
+        return row
+
+
+class ClsPreprocessor(ResponsePreprocessor):
+
+    def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        res = super().preprocess(row)
+        res['label'] = int(res['label'])
+        return res
+
+
+class AutoPreprocessor:
+
+    def __init__(self, *, columns: Optional[Dict[str, str]] = None, **kwargs) -> None:
+        self.columns = columns or {}
+        self.kwargs = kwargs
+
+    def _get_preprocessor(self, dataset: DATASET_TYPE) -> RowPreprocessor:
+        features = dataset.features
+        for key in ['conversation', 'conversations', 'messages']:
+            if key in features:
+                return MessagesPreprocessor(**self.kwargs)
+        if 'instruction' in features and 'input' in features:
+            return AlpacaPreprocessor(**self.kwargs)
+        return ResponsePreprocessor(**self.kwargs)
+
+    def __call__(
+        self,
+        dataset: DATASET_TYPE,
+        *,
+        num_proc: int = 1,
+        load_from_cache_file: bool = True,
+        **kwargs,
+    ) -> DATASET_TYPE:
+        dataset = RowPreprocessor.safe_rename_columns(dataset, self.columns)
+        preprocessor = self._get_preprocessor(dataset)
+        return preprocessor(dataset, num_proc=num_proc, load_from_cache_file=load_from_cache_file, **kwargs)

swift/dataset/preprocessor/extra.py

@@ -0,0 +1,111 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import numpy as np
+from typing import Any, Dict, List, Optional
+
+from .core import ResponsePreprocessor
+
+
+class GroundingMixin:
+    """This class offers prompts to the grounding task"""
+    task_type: Optional[str] = None
+
+    _grounding_language_mixin = [0.8, 0.2]
+    _grounding_prompts = {
+        'grounding': {
+            'en': [('<ref-object>', '<bbox>'), ('The positions of <ref-object> is', '<bbox>'),
+                   ('Find the positions of <ref-object>', '<bbox>'), ('Where is <ref-object>', '<bbox>'),
+                   ('Find <ref-object>', '<bbox>'), ('Show me <ref-object>', '<bbox>'),
+                   ('Detect <ref-object>', '<bbox>'), ('Locate <ref-object>', '<bbox>'),
+                   ('Tell me the location of <ref-object>', '<bbox>'), ('Give the location of <ref-object>', '<bbox>'),
+                   ('Provide the bounding box coordinate of <ref-object>', '<bbox>')],
+            'zh': [('<ref-object>', '<bbox>'), ('<ref-object>的位置在图片中', '<bbox>'), ('<ref-object>在图片中', '<bbox>'),
+                   ('<ref-object>在', '<bbox>'), ('找到<ref-object>的位置', '<bbox>'), ('<ref-object>在哪里', '<bbox>'),
+                   ('提供<ref-object>的坐标位置', '<bbox>')]
+        },
+        'caption': {
+            'en': [
+                ('<bbox>', '<ref-object>'),
+                ('The object at position <bbox>', '<ref-object>'),
+                ('This <bbox> is', '<ref-object>'),
+                ('What is the object at <bbox>', '<ref-object>'),
+                ('Describe <bbox>', '<ref-object>'),
+                ('<bbox> is', '<ref-object>'),
+                ('The bounding box coordinate <bbox> contains', '<ref-object>'),
+            ],
+            'zh': [
+                ('<bbox>', '<ref-object>'),
+                ('<bbox>是什么', '<ref-object>'),
+                ('<bbox>的位置包含', '<ref-object>'),
+                ('描述<bbox>', '<ref-object>'),
+                ('<bbox>中是', '<ref-object>'),
+                ('坐标<bbox>描述了什么', '<ref-object>'),
+                ('描述<bbox>中的事物', '<ref-object>'),
+            ]
+        },
+    }
+
+    def construct_grounding_prompt(self):
+        # TODO Only support one bbox to one object
+        lang = np.random.choice(['en', 'zh'], p=[0.8, 0.2])
+        prompts = GroundingMixin._grounding_prompts[self.task_type][lang]
+        query, response = prompts[np.random.choice(range(len(prompts)))]
+        return query, response
+
+
+class TextGenerationPreprocessor(ResponsePreprocessor):
+
+    def __init__(self,
+                 *,
+                 prompt: str,
+                 query_tag: str = '{{QUERY}}',
+                 columns: Optional[Dict[str, str]] = None,
+                 **kwargs) -> None:
+        self.query_tag = query_tag
+        self.prompt = prompt
+        super().__init__(columns=columns, **kwargs)
+
+    def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]:
+        row['query'] = self.prompt.replace(self.query_tag, row['query'])
+        return super().preprocess(row)
+
+
+class ClsGenerationPreprocessor(ResponsePreprocessor):
+
+    def __init__(self,
+                 labels: List[str],
+                 *,
+                 task: str,
+                 is_pair_seq: bool = False,
+                 columns: Optional[Dict[str, str]] = None,
+                 **kwargs) -> None:
+        self.labels = labels
+        self.task = task
+        self.is_pair_seq = is_pair_seq
+
+        category = ', '.join(labels)
+        self.sentence2_key = 'sentence2'
+        self.label_key = 'label'
+        if is_pair_seq:
+            self.sentence_key = 'sentence1'
+            inputs = 'Sentence1: {sentence1}\nSentence2: {sentence2}'
+        else:
+            self.sentence_key = 'sentence'
+            inputs = 'Sentence: {sentence}'
+        self.prompt = f"""Task: {task}
+{inputs}
+Category: {category}
+Output:"""
+        super().__init__(columns=columns, **kwargs)
+
+    def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        label = row.pop(self.label_key, None)
+        if label is None:
+            return
+
+        if self.is_pair_seq:
+            query = self.prompt.format(sentence1=row.pop(self.sentence_key), sentence2=row.pop(self.sentence2_key))
+        else:
+            query = self.prompt.format(sentence=row.pop(self.sentence_key))
+        row['query'] = query
+        row['response'] = self.labels[int(label)]
+        return super().preprocess(row)

swift/dataset/register.py

@@ -0,0 +1,115 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import json
+import os
+from copy import deepcopy
+from typing import Any, Dict, List, Optional, Union
+
+from swift.utils import get_logger, use_hf_hub
+from .dataset_meta import DATASET_MAPPING, DatasetMeta, SubsetDataset
+from .preprocessor import AutoPreprocessor, MessagesPreprocessor
+
+logger = get_logger()
+
+
+def get_dataset_list():
+    datasets = []
+    for key in DATASET_MAPPING:
+        if use_hf_hub():
+            if key[1]:
+                datasets.append(key[1])
+        else:
+            if key[0]:
+                datasets.append(key[0])
+    return datasets
+
+
+def register_dataset(dataset_meta: DatasetMeta, *, exist_ok: bool = False) -> None:
+    """Register dataset
+
+    Args:
+        dataset_meta: The `DatasetMeta` info of the dataset.
+        exist_ok: If the dataset id exists, raise error or update it.
+    """
+    if dataset_meta.dataset_name:
+        dataset_name = dataset_meta.dataset_name
+    else:
+        dataset_name = dataset_meta.ms_dataset_id, dataset_meta.hf_dataset_id, dataset_meta.dataset_path
+    if not exist_ok and dataset_name in DATASET_MAPPING:
+        raise ValueError(f'The `{dataset_name}` has already been registered in the DATASET_MAPPING.')
+
+    DATASET_MAPPING[dataset_name] = dataset_meta
+
+
+def _preprocess_d_info(d_info: Dict[str, Any], *, base_dir: Optional[str] = None) -> Dict[str, Any]:
+    d_info = deepcopy(d_info)
+
+    columns = None
+    if 'columns' in d_info:
+        columns = d_info.pop('columns')
+
+    if 'messages' in d_info:
+        d_info['preprocess_func'] = MessagesPreprocessor(**d_info.pop('messages'), columns=columns)
+    else:
+        d_info['preprocess_func'] = AutoPreprocessor(columns=columns)
+
+    if 'dataset_path' in d_info:
+        dataset_path = d_info.pop('dataset_path')
+        if base_dir is not None and not os.path.isabs(dataset_path):
+            dataset_path = os.path.join(base_dir, dataset_path)
+        dataset_path = os.path.abspath(os.path.expanduser(dataset_path))
+
+        d_info['dataset_path'] = dataset_path
+
+    if 'subsets' in d_info:
+        subsets = d_info.pop('subsets')
+        for i, subset in enumerate(subsets):
+            if isinstance(subset, dict):
+                subsets[i] = SubsetDataset(**_preprocess_d_info(subset))
+        d_info['subsets'] = subsets
+    return d_info
+
+
+def _register_d_info(d_info: Dict[str, Any], *, base_dir: Optional[str] = None) -> DatasetMeta:
+    """Register a single dataset to dataset mapping
+
+    Args:
+        d_info: The dataset info
+    """
+    d_info = _preprocess_d_info(d_info, base_dir=base_dir)
+    dataset_meta = DatasetMeta(**d_info)
+    register_dataset(dataset_meta)
+    return dataset_meta
+
+
+def register_dataset_info(dataset_info: Union[str, List[str], None] = None) -> List[DatasetMeta]:
+    """Register dataset from the `dataset_info.json` or a custom dataset info file
+    This is used to deal with the datasets defined in the json info file.
+
+    Args:
+        dataset_info: The dataset info path
+    """
+    # dataset_info_path: path, json or None
+    if dataset_info is None:
+        dataset_info = os.path.join(os.path.dirname(__file__), 'data', 'dataset_info.json')
+    assert isinstance(dataset_info, (str, list))
+    base_dir = None
+    log_msg = None
+    if isinstance(dataset_info, str):
+        dataset_path = os.path.abspath(os.path.expanduser(dataset_info))
+        if os.path.isfile(dataset_path):
+            log_msg = dataset_path
+            base_dir = os.path.dirname(dataset_path)
+            with open(dataset_path, 'r', encoding='utf-8') as f:
+                dataset_info = json.load(f)
+        else:
+            dataset_info = json.loads(dataset_info)  # json
+    if len(dataset_info) == 0:
+        return []
+    res = []
+    for d_info in dataset_info:
+        res.append(_register_d_info(d_info, base_dir=base_dir))
+
+    if log_msg is None:
+        log_msg = dataset_info if len(dataset_info) < 5 else list(dataset_info.keys())
+    logger.info(f'Successfully registered `{log_msg}`.')
+    return res

swift/dataset/utils.py

@@ -0,0 +1,153 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import inspect
+import numpy as np
+import os
+import tempfile
+from datasets import Dataset as HfDataset
+from modelscope.hub.utils.utils import get_cache_dir
+from torch.utils.data import Dataset
+from typing import Any, Callable, Dict, Optional, Union
+
+from swift.template import MaxLengthError, Template
+from swift.utils import get_logger
+from .preprocessor import RowPreprocessor
+
+logger = get_logger()
+
+
+def sample_dataset(
+        dataset: HfDataset,
+        dataset_sample: Optional[int],
+        shuffle: bool = True,
+        random_state: Optional[np.random.RandomState] = None,
+        shuffle_all: bool = False,  # For compatibility, this defaults to False.
+) -> HfDataset:
+    """Sample dataset by a dataset_sample number
+    Args:
+        dataset: The dataset instance, iterable dataset is not supported
+        dataset_sample: The sample number
+        shuffle: Whether to perform random sampling on non-streaming datasets
+        random_state: The random state
+    Returns:
+        The sampled dataset
+    """
+    if dataset_sample is None:
+        return dataset
+
+    n_repeat_sample = dataset_sample // len(dataset)
+    n_remain_sample = dataset_sample % len(dataset)
+    if n_repeat_sample >= 1 and n_remain_sample >= 1:
+        logger.warning(f'dataset_sample:{dataset_sample} is greater than len(dataset):{len(dataset)}, '
+                       'repeated sampling will be performed.')
+    idx = np.tile(range(len(dataset)), n_repeat_sample)
+    if random_state is None:
+        random_state = np.random.RandomState()
+    if n_remain_sample >= 1:
+        if shuffle:
+            idx_remain = random_state.permutation(len(dataset))[:n_remain_sample]
+        else:
+            idx_remain = np.arange(n_remain_sample)
+        idx = np.concatenate([idx, idx_remain])
+    if n_repeat_sample >= 1 and shuffle and shuffle_all:
+        random_state.shuffle(idx)
+    dataset = dataset.select(idx)
+    return dataset
+
+
+class LazyLLMDataset(Dataset):
+    """This class if used to lazy tokenize the dataset, and skips bad ones when training"""
+
+    def __init__(self,
+                 dataset: HfDataset,
+                 encode_func: Callable[[Dict[str, Any]], Dict[str, Any]],
+                 *,
+                 n_try_fetch: int = 10,
+                 strict: bool = False,
+                 random_state: Optional[Union[np.random.RandomState, int]] = None,
+                 traceback_limit: int = 10) -> None:
+        self.dataset = dataset
+        self.encode_func = encode_func
+
+        n_try_fetch = 1 if strict else min(n_try_fetch, len(self.dataset))
+        assert n_try_fetch >= 1
+        self.strict = strict
+        self.n_try_fetch = n_try_fetch
+
+        if not isinstance(random_state, np.random.RandomState):
+            random_state = np.random.RandomState(random_state)
+        self.random_state = random_state
+
+        self.traceback_limit = traceback_limit
+        self._traceback_counter = 0
+        self._idx = 0
+        self._idx_list = self.random_state.permutation(len(self.dataset)).tolist()
+
+    def __getitem__(self, idx: int) -> Dict[str, Any]:
+        if isinstance(idx, str):
+            return self.dataset[idx]
+        for i in range(self.n_try_fetch):
+            if i > 0:
+                idx = self._idx_list[self._idx]
+                self._idx = (self._idx + 1) % len(self.dataset)
+            data = self.dataset[idx]
+            try:
+                return self.encode_func(data, return_length=True)
+            except Exception as e:
+                if self.strict:
+                    logger.warning('To avoid errors, you can pass `strict=False`.')
+                    raise
+                if isinstance(e, MaxLengthError):
+                    continue
+                if self.traceback_limit is not None and self._traceback_counter < self.traceback_limit:
+                    import traceback
+                    logger.info(traceback.format_exc())
+                    logger.warning('👆👆👆There are errors in the template.encode, '
+                                   'and another piece of data will be randomly selected.')
+                    self._traceback_counter += 1
+
+        raise ValueError('Failed to retrieve the dataset. You can avoid this issue by increasing `max_length` or '
+                         'modifying the `truncation_strategy`.')
+
+    def __len__(self) -> int:
+        return len(self.dataset)
+
+
+class EncodePreprocessor(RowPreprocessor):
+
+    def __init__(self, template: 'Template'):
+        super().__init__()
+        self.template = template
+
+    def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        return self.template.encode(row, return_length=True)
+
+
+class AddLengthPreprocessor(EncodePreprocessor):
+
+    def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        encoded = super().preprocess(row)
+        row['lengths'] = encoded['lengths']
+        return row
+
+
+TEMP_DIR_POOL = {}
+
+
+def get_temporary_cache_files_directory(prefix=None):
+    if prefix is None:
+        import datasets.config
+        prefix = datasets.config.TEMP_CACHE_DIR_PREFIX
+    if prefix in TEMP_DIR_POOL:
+        TEMP_DIR = TEMP_DIR_POOL[prefix]
+    else:
+        tmp_dir = os.path.join(get_cache_dir(), 'tmp')
+        os.makedirs(tmp_dir, exist_ok=True)
+        kwargs = {}
+        parameters = inspect.signature(tempfile.TemporaryDirectory.__init__).parameters
+        if 'ignore_cleanup_errors' in parameters:
+            kwargs['ignore_cleanup_errors'] = True
+        TEMP_DIR = tempfile.TemporaryDirectory(prefix=prefix, dir=tmp_dir, **kwargs)
+        logger.info(f'create tmp_dir: {TEMP_DIR.name}')
+        TEMP_DIR_POOL[prefix] = TEMP_DIR
+
+    return TEMP_DIR.name

swift/hub/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from .hub import HFHub, MSHub, get_hub