chore: import upstream snapshot with attribution

swift/loss_scale/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from .base import ALL_BASE_STRATEGY, ConfigLossScale, LossScale
+from .mapping import get_loss_scale, loss_scale_map

swift/loss_scale/agent.py

@@ -0,0 +1,32 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+
+from typing import Optional
+
+from .base import ConfigLossScale
+from .utils import calculate_loss_scale
+
+
+class AgentFlanLossScale(ConfigLossScale):
+    is_binary = False
+    loss_scale_config = 'agentflan.json'
+
+    def get_loss_scale(self, context: str, *, query: Optional[str] = None, **kwargs):
+        if isinstance(context, str):
+            return calculate_loss_scale(query, context, self.loss_scale_map['response'], self.loss_scale_map['query'])
+        return super().get_loss_scale(context)
+
+
+class REACTLossScale(ConfigLossScale):
+    loss_scale_config = 'react.json'
+
+
+class QwenLossScale(ConfigLossScale):
+    loss_scale_config = 'qwen.json'
+
+
+class HermesLossScale(ConfigLossScale):
+    loss_scale_config = 'hermes.json'
+
+
+class AlphaUmiLossScale(ConfigLossScale):
+    loss_scale_config = 'alpha_umi.json'

swift/loss_scale/base.py

@@ -0,0 +1,236 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import json
+import os
+from typing import List, Literal, Optional, Tuple
+
+from swift.template import ContextType, Messages, get_last_user_round
+from .utils import calculate_loss_scale
+
+ALL_BASE_STRATEGY = ['default', 'last_round', 'all']
+
+
+class LossScale:
+    """Base class for loss scaling in training.
+
+        This class provides a flexible framework for controlling loss computation weights
+        across different parts of the context (e.g., system prompts, user queries, assistant
+        responses) during model training. Different strategies can be applied to selectively
+        train on specific portions.
+
+        Attributes:
+            is_binary (bool, optional): Indicates whether loss_scale contains only 0 and 1.
+                If True, loss_scale will be replaced by labels to stay compatible with
+                acceleration techniques such as liger_kernel.
+                If False, an additional 'loss_scale' key will be stored and the
+                corresponding loss function will be used.
+            base_strategy (str): Base strategy for loss computation. One of 'default',
+                'last_round', or 'all'.
+                - 'default': Only compute loss on assistant responses
+                - 'last_round': Only compute loss on the last round's assistant response
+                - 'all': Compute loss on all parts
+    """
+    is_binary = True
+
+    def __init__(self, base_strategy: Literal['default', 'last_round', 'all'] = 'default'):
+        """Initialize the loss scale object.
+
+        Args:
+            base_strategy: Base strategy for loss computation. One of 'default',
+                'last_round', or 'all'. Defaults to 'default'.
+
+        Raises:
+            ValueError: If the provided base_strategy is not in the allowed list.
+        """
+        if base_strategy not in ALL_BASE_STRATEGY:
+            raise ValueError(f'ALL_BASE_STRATEGY: {ALL_BASE_STRATEGY}, base_strategy: {base_strategy}')
+        self.base_strategy = base_strategy
+
+    def get_loss_scale(self, context: str, **kwargs) -> Tuple[List[str], List[float]]:
+        """Calculate loss scale for the given context.
+
+        This is a base implementation that subclasses can override to implement
+        custom loss scaling logic.
+
+        Args:
+            context: The input context (string).
+            **kwargs: Additional keyword arguments, such as query (the query of the
+                current round).
+
+        Returns:
+            A tuple containing:
+                - List[str]: List of contexts, potentially split into multiple parts
+                - List[float]: Corresponding loss scale values, one-to-one with contexts
+        """
+        return [context], [1.]
+
+    def __call__(self, context_list: List[str], context_types: List[ContextType], messages: Messages,
+                 **kwargs) -> Tuple[List[str], List[float]]:
+        """Process the complete conversation context and return contexts with loss scales.
+
+            This method iterates through all context segments and determines the loss scale
+            for each based on the context type and base strategy. It handles special cases
+            such as explicitly specified loss values in messages and pre-computed loss scales.
+
+            Args:
+                context_list: List of context strings or dicts, each representing a segment
+                    of the conversation.
+                context_types: List of context types corresponding to each context, indicating
+                    whether it's a system prompt, user query, assistant response, etc.
+                messages: Complete message list containing the conversation history.
+
+            Returns:
+                A tuple containing:
+                    - List[str]: Processed context list, potentially expanded if contexts
+                        are split into multiple parts
+                    - List[float]: Loss scale values corresponding one-to-one with the
+                        returned context list
+        """
+        res_context_list = []
+        res_loss_scale = []
+        i = 0
+        last_user_round = get_last_user_round(messages)
+        for context, context_type in zip(context_list, context_types):
+            is_last_round = 2 * i >= last_user_round
+            query, loss, loss_scale = None, None, None
+            if context_type == ContextType.RESPONSE:
+                query = messages[2 * i]['content']
+                # Currently, we only support applying loss/mask to the response part.
+                loss = messages[2 * i + 1].get('loss')
+                loss_scale = messages[2 * i + 1].get('loss_scale')
+                assert context == messages[2 * i + 1]['content']
+                i += 1
+            if not isinstance(context, list) or (len(context) > 0 and isinstance(context[0], int)):
+                context = [context]
+            for j in range(len(context)):
+                new_context, new_loss_scale = self._inner_call(
+                    context[j], context_type, query, loss[j] if isinstance(loss, list) else loss,
+                    loss_scale[j] if isinstance(loss_scale, list) else loss_scale, is_last_round)
+                res_context_list += new_context
+                res_loss_scale += new_loss_scale
+
+        # The values in loss_scale_list correspond one-to-one with the values in context_list.
+        return res_context_list, res_loss_scale
+
+    def _inner_call(self, context, context_type, query, loss, loss_scale, is_last_round):
+        if isinstance(context, dict) and 'loss_scale' in context:
+            new_context = [[token] for token in context['token_ids']]
+            loss_scale = context['loss_scale']
+        else:
+            if isinstance(context, dict) and 'token_ids' in context:
+                context = context['token_ids']
+            is_assistant = context_type in {ContextType.RESPONSE, ContextType.SUFFIX}
+            if loss or loss is None and (self.base_strategy == 'all' or
+                                         (self.base_strategy == 'default' and is_assistant) or
+                                         (self.base_strategy == 'last_round' and is_assistant and is_last_round)):
+                if loss_scale is None:
+                    new_context, loss_scale = self.get_loss_scale(context, query=query)
+                else:
+                    new_context, loss_scale = [context], [loss_scale]
+            else:
+                new_context, loss_scale = [context], [0.]
+        return new_context, loss_scale
+
+    @property
+    def is_binary_loss_scale(self):
+        """Check if loss scale values are binary (only 0 and 1)."""
+        return self.is_binary
+
+
+class ConfigLossScale(LossScale):
+    """Loss scale class that loads configuration from a JSON file.
+
+    This class extends LossScale to support loading predefined loss scale mappings
+    from a configuration file. The mappings can specify different loss weights for
+    different tokens or segments based on the content.
+
+    Attributes:
+        loss_scale_config (str, optional): Path to the loss scale configuration file
+            relative to the 'config' directory.
+        loss_scale_map (dict, optional): Dictionary mapping loaded from the config file,
+            containing predefined loss scale values for specific patterns or tokens.
+    """
+    is_binary = None
+    loss_scale_config = None  # path
+
+    def __init__(self, base_strategy: Literal['default', 'last_round', 'all'] = 'default'):
+        """Initialize the config-based loss scale object.
+
+        Loads the loss scale configuration from a JSON file if loss_scale_config
+        is specified.
+
+        Args:
+            base_strategy: Base strategy for loss computation. One of 'default',
+                'last_round', or 'all'. Defaults to 'default'.
+        """
+        super().__init__(base_strategy)
+        self.loss_scale_map = None
+        if self.loss_scale_config is not None:
+            path = os.path.dirname(os.path.abspath(__file__))
+            config_path = os.path.join(path, 'config', self.loss_scale_config)
+            with open(config_path, 'r', encoding='utf-8') as json_file:
+                self.loss_scale_map = json.load(json_file)
+
+    @property
+    def is_binary_loss_scale(self):
+        if self.is_binary is not None:
+            return self.is_binary
+        if self.loss_scale_map is None:
+            return True
+        return all(scale in {0.0, 1.0} for lst in self.loss_scale_map.values() for scale in lst)
+
+    def get_loss_scale(self, context: str, *, query: Optional[str] = None, **kwargs):
+        """Calculate loss scale using the loaded configuration.
+
+        If context is a string, uses the configuration map to calculate loss scales
+        based on the query and context. Otherwise, falls back to the parent class
+        implementation.
+
+        Args:
+            context: The input context string.
+            query: The user query for the current round, used to determine
+                appropriate loss scaling based on the configuration.
+
+        Returns:
+            Tuple[List[str], List[float]]: List of context segments and their
+                corresponding loss scale values.
+        """
+        if isinstance(context, str):
+            return calculate_loss_scale(query, context, self.loss_scale_map)
+        return super().get_loss_scale(context)
+
+
+class ConcatLossScale(LossScale):
+    """Apply multiple loss scales sequentially.
+
+    The output segments of each underlying loss scale are fed into the next one,
+    and the corresponding weights are multiplied together. This makes it possible
+    to compose strategies such as ``hermes+ignore_empty_think``.
+    """
+
+    is_binary = None
+
+    def __init__(self,
+                 loss_scales: List[LossScale],
+                 base_strategy: Literal['default', 'last_round', 'all'] = 'default'):
+        super().__init__(base_strategy)
+        assert loss_scales, 'loss_scales must be a non-empty list'
+        self.loss_scales = loss_scales
+
+    def get_loss_scale(self, context, **kwargs):
+        contexts = [context]
+        weights = [1.0]
+        for ls in self.loss_scales:
+            new_contexts: List = []
+            new_weights: List[float] = []
+            for c, w in zip(contexts, weights):
+                sub_contexts, sub_weights = ls.get_loss_scale(c, **kwargs)
+                for sc, sw in zip(sub_contexts, sub_weights):
+                    new_contexts.append(sc)
+                    new_weights.append(w * sw)
+            contexts = new_contexts
+            weights = new_weights
+        return contexts, weights
+
+    @property
+    def is_binary_loss_scale(self):
+        return all(ls.is_binary_loss_scale for ls in self.loss_scales)

swift/loss_scale/config/agentflan.json

@@ -0,0 +1,22 @@
+{
+    "response":{
+        "Name:": [1.0, 3.0],
+        "Action:": [1.0, 3.0],
+        "ACTION:": [1.0,3.0],
+        "Tool:": [1.0, 3.0],
+        "Command": [1.0, 3.0],
+        "Arguments:": [1.0, 3.0],
+        "action input": [1.0, 3.0],
+        "ACTION_INPUT:":[1.0, 3.0],
+        "Action Input:": [1.0, 3.0],
+        "Thought:": [1.0, 1.0],
+        "Final Answer:": [1.0, 1.0],
+        "Observation:": [2.0, 0.0]
+    },
+    "query":{
+        "What is the tool you want to use": [3.0],
+        "What are the required parameter names": [3.0],
+        "What is the value of": [3.0],
+        "What are the required parameter names for this tool": [3.0]
+    }
+}

swift/loss_scale/config/alpha_umi.json

@@ -0,0 +1,8 @@
+{
+    "Action:": [2.0, 2.0],
+    "Action Input:": [2.0, 2.0],
+    "Thought:": [1.0, 1.0],
+    "Final Answer:": [1.0, 1.0],
+    "Observation:": [2.0, 0.0],
+    "Next:": [2,0, 2.0]
+}

swift/loss_scale/config/hermes.json

@@ -0,0 +1,3 @@
+{
+    "<tool_call>.+?</tool_call>": [2.0]
+}

swift/loss_scale/config/ignore_empty_think.json

@@ -0,0 +1,7 @@
+{
+    "^<think>\\s*</think>\\s*": [0.0],
+    "^<seed:think><seed:cot_budget_reflect>The current thinking budget is 0, so I will directly start answering the question.</seed:cot_budget_reflect>\\n</seed:think>\\s*": [0.0],
+    "^</think>\\s*": [0.0],
+    "^<\\|channel>thought\\n<channel\\|>": [0.0],
+    "^</mm:think>\\s*": [0.0]
+}

swift/loss_scale/config/qwen.json

@@ -0,0 +1,6 @@
+{
+    "✿FUNCTION✿:": [2.0, 2.0],
+    "✿ARGS✿:": [2.0, 2.0],
+    "✿RETURN✿:": [1.0, 1.0],
+    "✿RESULT✿:": [2.0, 0.0]
+}

swift/loss_scale/config/react.json

@@ -0,0 +1,7 @@
+{
+    "Action:": [2.0, 2.0],
+    "Action Input:": [2.0, 2.0],
+    "Thought:": [1.0, 1.0],
+    "Final Answer:": [1.0, 1.0],
+    "Observation:": [2.0, 0.0]
+}

swift/loss_scale/mapping.py

@@ -0,0 +1,57 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from .agent import AgentFlanLossScale, AlphaUmiLossScale, HermesLossScale, QwenLossScale, REACTLossScale
+from .base import ALL_BASE_STRATEGY, ConcatLossScale, LossScale
+from .other import IgnoreEmptyThinkLossScale
+
+# Add your loss scale here, use --loss_scale xxx to train
+loss_scale_map = {
+    'base': LossScale,
+    'ignore_empty_think': IgnoreEmptyThinkLossScale,
+    # agent
+    'react': REACTLossScale,
+    'hermes': HermesLossScale,
+    'qwen': QwenLossScale,
+    'agentflan': AgentFlanLossScale,
+    'alpha_umi': AlphaUmiLossScale,
+}
+
+
+def get_loss_scale(loss_scale: str) -> LossScale:
+    """Factory function to create a loss scale object from a string specification.
+
+    The loss_scale string supports the following formats (segments separated by '+'):
+    1. A strategy name alone (e.g., 'default', 'last_round', 'all') - uses base LossScale
+    2. A loss scale type alone (e.g., 'hermes', 'react') - uses 'default' strategy
+    3. A strategy name followed by a loss scale type (e.g., 'default+react', 'last_round+qwen')
+    4. Multiple loss scale types chained together, optionally led by a base strategy
+       (e.g., 'hermes+ignore_empty_think', 'last_round+hermes+ignore_empty_think').
+       The chained loss scales are applied sequentially: each loss scale processes the
+       output of the previous one and the corresponding weights are multiplied together.
+
+    Args:
+        loss_scale: String specifying the loss scale configuration.
+
+    Returns:
+        LossScale: An instance of the appropriate LossScale subclass. When multiple loss
+            scale types are specified, a ``ConcatLossScale`` wrapping them is returned.
+
+    Examples:
+        >>> get_loss_scale('default')  # Uses default strategy with base LossScale
+        >>> get_loss_scale('react')  # Uses default strategy with REACTLossScale
+        >>> get_loss_scale('last_round+hermes')  # last_round strategy with HermesLossScale
+        >>> get_loss_scale('last_round+hermes+ignore_empty_think')  # chain hermes then ignore_empty_think
+    """
+    parts = loss_scale.split('+')
+    if parts[0] in ALL_BASE_STRATEGY:
+        base_strategy = parts[0]
+        ls_names = parts[1:] or ['base']
+    else:
+        base_strategy = 'default'
+        ls_names = parts
+    if len(ls_names) == 1:
+        return loss_scale_map[ls_names[0]](base_strategy)
+    # The base_strategy is owned by the outer ConcatLossScale; sub loss scales only
+    # contribute their `get_loss_scale` (which does not reference base_strategy), so
+    # any valid placeholder ('default') is fine here.
+    sub_loss_scales = [loss_scale_map[name]('default') for name in ls_names]
+    return ConcatLossScale(sub_loss_scales, base_strategy)

swift/loss_scale/other.py

@@ -0,0 +1,6 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from .base import ConfigLossScale
+
+
+class IgnoreEmptyThinkLossScale(ConfigLossScale):
+    loss_scale_config = 'ignore_empty_think.json'

swift/loss_scale/utils.py

@@ -0,0 +1,59 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from typing import Dict, List, Optional, Tuple
+
+from swift.template import split_str_parts_by
+
+
+def calculate_loss_scale(query: str,
+                         response: str,
+                         response_loss_scale_map: Dict[str, list],
+                         query_loss_scale_map: Optional[Dict[str, list]] = None) -> Tuple[List[str], List[float]]:
+    """Calculate the loss scale by splitting the agent response.
+
+    This algorithm comes from paper: https://arxiv.org/pdf/2309.00986.pdf
+
+    Agent response format:
+
+    ```text
+        Thought: you should always think about what to do
+        Action: the action to take, should be one of the above tools[fire_recognition,
+            fire_alert, call_police, call_fireman]
+        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
+    ```
+    Returns:
+        A tuple of agent response parts and their weights.
+    """
+    # query loss scale map
+    if query_loss_scale_map is not None:
+        for key in query_loss_scale_map.keys():
+            if key in query:
+                if isinstance(query_loss_scale_map[key], (float, int)):
+                    query_loss_scale_map[key] = [query_loss_scale_map[key]]
+                loss_scale_value = query_loss_scale_map[key][0]
+                return [response], [float(loss_scale_value)]
+    delimiters = [k for k, v in response_loss_scale_map.items() if len(v) == 2]
+    if delimiters:
+        agent_parts = split_str_parts_by(response, delimiters)
+    else:
+        regex_delimiters = [k for k, v in response_loss_scale_map.items() if len(v) == 1]
+        agent_parts = split_str_parts_by(response, regex_delimiters, regex_mode=True)
+    weights = []
+    agent_content = []
+    for c in agent_parts:
+        if c['key'] in response_loss_scale_map:
+            loss_scale = response_loss_scale_map[c['key']]
+            assert len(loss_scale) in {1, 2}, f'loss_scale: {loss_scale}'
+            if len(loss_scale) == 1:
+                weights += loss_scale
+                agent_content.append(c['content'])
+            else:
+                weights += loss_scale
+                agent_content += [c['key'], c['content']]
+        else:
+            weights.append(1.)
+            agent_content.append(c['content'])
+    return agent_content, weights