chore: import upstream snapshot with attribution

swift/pipelines/sampling/__init__.py

@@ -0,0 +1 @@
+from .sampling import sampling_main

swift/pipelines/sampling/base.py

@@ -0,0 +1,59 @@
+from typing import Any, Dict, List
+
+from swift.arguments import SamplingArguments
+from swift.infer_engine import TransformersEngine
+from swift.ray_utils import RayHelper
+from swift.rewards import orms, prms
+from swift.utils import get_logger
+
+logger = get_logger()
+
+
+class Sampler:
+
+    def __init__(self, input_args: SamplingArguments):
+        self.args = input_args
+        self.template = None
+        self.processor = None
+        self.prm_model = None
+        self.orm_model = None
+        self._prepare_model_tokenizer()
+        self._prepare_template()
+        self._prepare_prm()
+        self._prepare_orm()
+
+    def _prepare_model_tokenizer(self):
+        args = self.args
+        _, self.processor = args.get_model_processor(load_model=False)
+
+    @RayHelper.function(group='prm')
+    def _prepare_prm(self):
+        if self.args.prm_model is None:
+            self.prm_model = None
+            logger.warning('prm_model is None.')
+        elif self.args.prm_model in prms:
+            self.prm_model = prms[self.args.prm_model]()
+        else:
+            self.prm_model = TransformersEngine(self.args.prm_model, max_batch_size=64)
+
+    @RayHelper.function(group='orm')
+    def _prepare_orm(self):
+        if self.args.orm_model is None:
+            self.orm_model = None
+            logger.warning('orm_model is None.')
+        elif self.args.orm_model in orms:
+            self.orm_model = orms[self.args.orm_model]()
+        else:
+            self.orm_model = TransformersEngine(self.args.orm_model, max_batch_size=64)
+
+    def _prepare_template(self) -> None:
+        template = self.args.get_template(self.processor)
+        self.template = template
+        self.template.set_mode('train')
+
+    def truncate_input(self, slices: List[Dict[str, Any]]):
+        """Truncate the input rows to avoid hitting the max length of the policy model"""
+        return slices
+
+    def do_sample(self, data):
+        raise NotImplementedError

swift/pipelines/sampling/distill_sampler.py

@@ -0,0 +1,160 @@
+import os
+from copy import deepcopy
+from openai import OpenAI
+from typing import List, Optional
+
+from swift.infer_engine import InferRequest, RequestConfig
+from swift.ray_utils import RayHelper
+from .utils import get_messages_md5
+from .vanilla_sampler import VanillaSampler
+
+
+class OpenAIEngine:
+
+    def __init__(
+        self,
+        model: str,
+        stream: bool = False,
+        base_url: str = 'https://dashscope.aliyuncs.com/compatible-mode/v1',
+        api_key: str = '',
+        **kwargs,
+    ):
+        self.model = model
+        self.stream = stream
+        self.client = OpenAI(api_key=api_key if api_key else os.getenv('OPENAI_API_KEY'), base_url=base_url, **kwargs)
+
+    def infer(
+        self,
+        infer_requests: List[InferRequest],
+        request_config: Optional[RequestConfig] = None,
+    ):
+        resp_contents = []
+        for infer_request in infer_requests:
+            completion = self.client.chat.completions.create(
+                model=self.model,
+                messages=infer_request['messages'],
+                temperature=request_config.temperature,
+                top_p=request_config.top_p,
+                max_tokens=request_config.max_tokens,
+                stream=self.stream,
+            )
+            reasoning_content = None
+            if self.stream:
+                reasoning_content = ''
+                content = ''
+                for chunk in completion:
+                    chunk_choices = chunk.choices
+                    if len(chunk_choices) == 0:
+                        continue
+                    reasoning_chunk = chunk_choices[0].delta.reasoning_content if hasattr(
+                        chunk_choices[0].delta, 'reasoning_content') else ''
+                    answer_chunk = chunk_choices[0].delta.content
+                    if reasoning_chunk:
+                        reasoning_content += reasoning_chunk
+                    elif answer_chunk:
+                        content += answer_chunk
+            else:
+                if hasattr(completion.choices[0].message, 'reasoning_content'):
+                    reasoning_content = completion.choices[0].message.reasoning_content
+                content = completion.choices[0].message.content
+            assert len(content) > 0, 'Empty completion'
+            if reasoning_content:
+                resp_content = f'<think>{reasoning_content}</think>\n\n<answer>{content}</answer>'
+            else:
+                resp_content = content
+            resp_contents.append(resp_content)
+
+        return resp_contents
+
+
+@RayHelper.worker(group=['sampler', 'prm', 'orm'])
+class DistillSampler(VanillaSampler):
+
+    def __init__(self, *args, **kwargs):
+        super(VanillaSampler, self).__init__(*args, **kwargs)
+        assert self.args.sampler_engine == 'client'
+        self._prepare_sampler()
+        self.caches = self.read_cache()
+
+    @RayHelper.function(group='sampler')
+    def _prepare_sampler(self):
+        self.infer_engine = OpenAIEngine(model=self.args.model, stream=self.args.stream, **self.args.engine_kwargs)
+        self.infer_engine.strict = False
+
+    def _prepare_model_tokenizer(self):
+        pass
+
+    def _prepare_template(self):
+        pass
+
+    def extract_choice(self, resp):
+        message = resp.choices[0].message
+        if hasattr(message, 'reasoning_content'):
+            reps_content = f'<think>{message.reasoning_content}</think>\n\n<answer>{message.content}</answer>'
+        else:
+            reps_content = message.content
+        return reps_content
+
+    @RayHelper.function(
+        group='sampler',
+        dispatch=lambda n, i, data:
+        ([{
+            'messages': data['messages'][i * len(data['messages']) // n:(i + 1) * len(data['messages']) // n]
+        }], {}),
+        collect='flatten')
+    def generate(self, data):
+        resp_all = []
+        infer_requests = []
+        sent = 0
+        rows = self.convert_data_to_rows(data)
+        for idx, row in enumerate(rows):
+            row = deepcopy(row)
+            messages = row['messages']
+            uuid = get_messages_md5(row)
+            if uuid in self.caches:
+                choices = self.caches[uuid]['choices']
+                if len(choices) == self.args.num_return_sequences:
+                    continue
+            if self.args.system:
+                if messages[0]['role'] == 'system':
+                    messages[0]['content'] = self.args.system
+                else:
+                    messages.insert(0, {'role': 'system', 'content': self.args.system})
+            if messages[-1]['role'] == 'assistant':
+                messages = messages[:-1]
+
+            row['messages'] = messages
+            infer_request = row
+            for i in range(self.args.num_return_sequences):
+                infer_requests.append(deepcopy(infer_request))
+            sent += 1
+
+        request_config = RequestConfig(
+            max_tokens=self.args.max_new_tokens,
+            temperature=self.args.temperature,
+            top_k=self.args.top_k,
+            top_p=self.args.top_p,
+        )
+
+        resp_list = []
+        if len(infer_requests) > 0:
+            resp_list = self.infer_engine.infer(infer_requests, request_config=request_config)
+
+        _cur = 0
+        for idx, row in enumerate(rows):
+            row = deepcopy(row)
+            uuid = get_messages_md5(row)
+            if uuid in self.caches:
+                choices = self.caches[uuid]['choices']
+                if len(choices) == self.args.num_return_sequences:
+                    row['choices'] = choices
+                    resp_all.append(row)
+                    continue
+
+            resps = row
+            resps['choices'] = []
+            for j in range(self.args.num_return_sequences * _cur, self.args.num_return_sequences * (_cur + 1)):
+                resps['choices'].append(resp_list[j])
+            resp_all.append(resps)
+            _cur += 1
+        return resp_all

swift/pipelines/sampling/sampling.py

@@ -0,0 +1,104 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import json
+import os
+import shutil
+import time
+from typing import List, Optional, Union
+
+from swift.arguments import SamplingArguments
+from swift.dataset import load_dataset
+from swift.utils import get_logger
+from ..base import SwiftPipeline
+from .distill_sampler import DistillSampler
+from .vanilla_sampler import VanillaSampler
+
+logger = get_logger()
+
+
+class SwiftSampling(SwiftPipeline):
+    args_class = SamplingArguments
+    args: args_class
+
+    def __init__(self, args: Optional[Union[List[str], SamplingArguments]] = None) -> None:
+        super().__init__(args)
+        self.args.save_args()
+        os.makedirs(self.args.output_dir, exist_ok=True)
+        self.cur_piece = 0
+        self.total_piece = 1
+
+        if self.args.data_range:
+            self.cur_piece, self.total_piece = self.args.data_range
+
+        if self.args.sampler_type == 'sample':
+            self.sampler = VanillaSampler(self.args)
+        elif self.args.sampler_type == 'distill':
+            self.sampler = DistillSampler(self.args)
+        else:
+            raise ValueError(f'Unsupported sampler type: {self.args.sampler_type}')
+
+    def _get_dataset(self):
+        args = self.args
+        dataset_kwargs = args.get_dataset_kwargs()
+        sampling_dataset, _ = load_dataset(
+            args.dataset, split_dataset_ratio=0., shuffle=args.dataset_shuffle, **dataset_kwargs)
+        logger.info(f'Sampling_dataset: {sampling_dataset}')
+        dataset_len = len(sampling_dataset)
+        piece_len = dataset_len // self.total_piece
+        sampling_dataset = sampling_dataset.select(range(piece_len * self.cur_piece, piece_len * (self.cur_piece + 1)))
+        return sampling_dataset
+
+    def run(self):
+        os.makedirs(self.args.output_dir, exist_ok=True)
+        iter_file = os.path.join(self.args.output_dir, self.args.output_file)
+        resume_file = os.path.join(self.args.output_dir, self.args.output_file + '.resume')
+        tmp_file = os.path.join(self.args.output_dir, self.args.output_file + '.tmp')
+        ckpt_state_file = os.path.join(self.args.output_dir, 'ckpt_state.json')
+        if os.path.exists(iter_file) and not self.args.override_exist_file:
+            return
+
+        index_resume = -1
+        write_mode = 'w'
+        if self.args.resume:
+            write_mode = 'a'
+            if os.path.exists(resume_file):
+                shutil.copyfile(resume_file, tmp_file)
+
+            if os.path.exists(ckpt_state_file):
+                with open(ckpt_state_file, 'r', encoding='utf-8') as ckpt_state:
+                    data = json.load(ckpt_state)
+                    index_resume = data.get('index', -1)
+                    logger.info(f'Loaded index_resume: {index_resume}')
+        else:
+            if os.path.exists(tmp_file):
+                os.remove(tmp_file)
+
+        dataset = self._get_dataset()
+        dataset_len = len(dataset)
+        total_iters = int(dataset_len // self.args.num_sampling_batch_size)
+
+        if self.args.num_sampling_batches is None or self.args.num_sampling_batches > total_iters:
+            self.args.num_sampling_batches = total_iters
+
+        with open(tmp_file, write_mode) as f:
+            for _index in range(self.args.num_sampling_batches):
+                if _index <= index_resume:
+                    continue
+                logger.info(f' Sampling index:{_index}')
+                slices = dataset[self.args.num_sampling_batch_size * _index:self.args.num_sampling_batch_size
+                                 * (_index + 1)]
+                slices = self.sampler.truncate_input(slices)
+                generated = self.sampler.do_sample(slices)
+                f.writelines(generated)
+                f.flush()
+                shutil.copy(tmp_file, resume_file)
+                with open(ckpt_state_file, 'w') as ckpt_state:
+                    json.dump({'index': _index}, ckpt_state)
+
+        if os.path.exists(iter_file):
+            shutil.move(iter_file, iter_file + '.' + str(int(time.time())))
+        shutil.move(resume_file, iter_file)
+        logger.info(f'Sample file {iter_file} generated.')
+
+
+def sampling_main(args: Optional[Union[List[str], SamplingArguments]] = None):
+    return SwiftSampling(args).main()

swift/pipelines/sampling/utils.py

@@ -0,0 +1,80 @@
+import hashlib
+import inspect
+import json
+import numpy as np
+from copy import copy
+from typing import Any, Dict, List, Optional
+
+from swift.infer_engine import ChatCompletionResponse, InferEngine, InferRequest, RequestConfig
+from swift.utils import get_logger
+
+logger = get_logger()
+
+
+def get_messages_md5(row: Dict[str, Any]):
+    row = copy(row)
+    row.pop('choices', None)
+    serialized = json.dumps(row, sort_keys=True)
+    return hashlib.md5(serialized.encode('utf-8')).hexdigest()
+
+
+def get_reward(model: Any,
+               infer_requests: List[InferRequest],
+               request_config: RequestConfig = None,
+               ground_truths: List[str] = None,
+               threshold: Optional[float] = None):
+    """Get reward from an RM model.
+
+    Args:
+        model: The model instance or an RM evaluator
+        infer_requests: Infer requests sent to the model
+        request_config: Infer config
+        ground_truths: The ground truth list
+        threshold: An optional threshold to generate the mask
+
+    Returns:
+        Tuple
+        Index 0: The min-max normalized scores matched the infer_requests
+        Index 1: The mask filtered by the threshold
+    """
+    infer_func = model.infer if isinstance(model, InferEngine) else model.__call__
+    parameters = inspect.signature(infer_func).parameters
+    gt_param = {}
+    if 'ground_truths' in parameters:
+        gt_param = {'ground_truths': ground_truths}
+    if isinstance(infer_requests[0], dict):
+        infer_requests = [InferRequest(messages=req['messages']) for req in infer_requests]
+    rewards = infer_func(infer_requests, request_config=request_config, **gt_param)
+    if isinstance(rewards[0], ChatCompletionResponse):
+        print('reward:', rewards[0].choices[0].message.content)
+        if isinstance(rewards[0].choices[0].message.content, str):
+            rewards = [float(r.choices[0].message.content.strip('[]')) for r in rewards]
+        elif isinstance(rewards[0].choices[0].message.content, list):
+            rewards = [float(min(r.choices[0].message.content)) for r in rewards]
+        else:
+            rewards = [float(r.choices[0].message.content) for r in rewards]
+    arr = []
+    for reward in rewards:
+        if isinstance(reward, (list, tuple)):
+            arr.append(min(reward))
+        else:
+            arr.append(float(reward))
+
+    _mask = np.array([True] * len(arr))
+    if threshold is not None:
+        # > not >=, orm caller passes 0, which will cause error
+        _mask = np.array([a > threshold for a in arr])
+
+    def normalize(arr):
+        min_val = np.min(arr)
+        max_val = np.max(arr)
+        if min_val == max_val:
+            if min_val == 0:
+                constant_value = 0.0
+            else:
+                constant_value = min(1.0, min_val)
+            return np.full_like(arr, fill_value=constant_value, dtype=np.float64)
+        normalized = (arr - min_val) / (max_val - min_val + 1e-5)
+        return normalized
+
+    return normalize(arr), _mask

swift/pipelines/sampling/vanilla_sampler.py

@@ -0,0 +1,233 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import json
+import numpy as np
+import os
+from copy import deepcopy
+
+from swift.infer_engine import RequestConfig, TransformersEngine
+from swift.ray_utils import RayHelper
+from swift.utils import get_logger
+from .base import Sampler
+from .utils import get_messages_md5, get_reward
+
+logger = get_logger()
+
+
+@RayHelper.worker(group=['sampler', 'prm', 'orm'])
+class VanillaSampler(Sampler):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._prepare_sampler()
+        self.caches = self.read_cache()
+
+    @RayHelper.function(group='sampler')
+    def _prepare_sampler(self):
+        if self.args.sampler_engine == 'transformers':
+            _Engine = TransformersEngine
+        elif self.args.sampler_engine == 'vllm':
+            from swift.infer_engine import VllmEngine
+            _Engine = VllmEngine
+        elif self.args.sampler_engine == 'lmdeploy':
+            from swift.infer_engine import LmdeployEngine
+            _Engine = LmdeployEngine
+        elif self.args.sampler_engine == 'no':
+            _Engine = None
+        else:
+            raise ValueError(f'Cannot find engine name: {self.args.sampler_engine}')
+        self.infer_engine = None
+        if _Engine:
+            self.infer_engine = _Engine(
+                self.args.model, model_type=self.args.model_type, template=self.template, **self.args.engine_kwargs)
+            self.infer_engine.strict = False
+
+    @RayHelper.function(group='sampler')
+    def read_cache(self):
+        cache_files = self.args.cache_files
+        caches = {}
+        for file in cache_files:
+            if not os.path.exists(file):
+                logger.warning(f'Cache file does not exist: {file}')
+                continue
+            with open(file, 'r', encoding='utf-8') as f:
+                for line in f.readlines():
+                    line = line.strip()
+                    if not line:
+                        continue
+
+                    content = json.loads(line)
+                    uuid = content['id']
+                    messages = content['messages']
+                    if uuid not in caches:
+                        caches[uuid] = {'choices': []}
+                    assert messages[-1]['role'] == 'assistant'
+                    caches[uuid]['choices'].append(messages[-1]['content'])
+        return caches
+
+    @staticmethod
+    def convert_data_to_rows(data):
+        rows = []
+        key = list(data.keys())[0]
+        data_len = len(data[key])
+        for idx in range(data_len):
+            row = {key: data[key][idx] for key in data}
+            if row.get('images') and 'bytes' in row['images'][0]:
+                row['images'] = [img['path'] for img in row['images']]
+            rows.append(row)
+        VanillaSampler.check_row_valid(rows)
+        return rows
+
+    @staticmethod
+    def check_row_valid(rows):
+        for row in rows:
+            assert not row.get('images') or all([isinstance(img, str) and img for img in row['images']])
+            assert not row.get('videos') or all([isinstance(video, str) and video for video in row['videos']])
+            assert not row.get('audios') or all([isinstance(audio, str) and audio for audio in row['audios']])
+
+    @RayHelper.function(
+        group='sampler',
+        dispatch=lambda n, i, data:
+        ([{
+            'messages': data['messages'][i * len(data['messages']) // n:(i + 1) * len(data['messages']) // n]
+        }], {}),
+        collect='flatten')
+    def generate(self, data):
+        resp_all = []
+        infer_requests = []
+        sent = 0
+        rows = self.convert_data_to_rows(data)
+        for idx, row in enumerate(rows):
+            row = deepcopy(row)
+            messages = row['messages']
+            uuid = get_messages_md5(row)
+            if uuid in self.caches:
+                choices = self.caches[uuid]['choices']
+                if len(choices) == self.args.num_return_sequences:
+                    continue
+            if self.args.system:
+                if messages[0]['role'] == 'system':
+                    messages[0]['content'] = self.args.system
+                else:
+                    messages.insert(0, {'role': 'system', 'content': self.args.system})
+            if messages[-1]['role'] == 'assistant':
+                messages = messages[:-1]
+
+            row['messages'] = messages
+            infer_request = row
+            for i in range(self.args.num_return_sequences):
+                infer_requests.append(deepcopy(infer_request))
+            sent += 1
+
+        request_config = RequestConfig(
+            max_tokens=self.args.max_new_tokens,
+            temperature=self.args.temperature,
+            top_k=self.args.top_k,
+            top_p=self.args.top_p,
+        )
+
+        resp_list = []
+        if len(infer_requests) > 0:
+            resp_list = self.infer_engine.infer(infer_requests, request_config=request_config)
+
+        _cur = 0
+        for idx, row in enumerate(rows):
+            row = deepcopy(row)
+            uuid = get_messages_md5(row)
+            if uuid in self.caches:
+                choices = self.caches[uuid]['choices']
+                if len(choices) == self.args.num_return_sequences:
+                    row['choices'] = choices
+                    resp_all.append(row)
+                    continue
+
+            resps = row
+            resps['choices'] = []
+            for j in range(self.args.num_return_sequences * _cur, self.args.num_return_sequences * (_cur + 1)):
+                if not isinstance(resp_list[j], Exception):
+                    resps['choices'].append(resp_list[j].choices[0].message.content)
+            if resps['choices']:
+                resp_all.append(resps)
+            _cur += 1
+        return resp_all
+
+    @RayHelper.function(group='orm', dispatch='slice', collect='flatten')
+    def get_orm_score(self, infer_requests, ground_truth):
+        return get_reward(
+            self.orm_model, infer_requests, ground_truths=[ground_truth] * len(infer_requests), threshold=0.0)
+
+    @RayHelper.function(group='prm', dispatch='slice', collect='flatten')
+    def get_prm_score(self, infer_requests, ground_truth):
+        return get_reward(
+            self.prm_model,
+            infer_requests,
+            ground_truths=[ground_truth] * len(infer_requests),
+            threshold=self.args.prm_threshold)
+
+    def do_sample(self, data):
+        generated = []
+        resp_all = self.generate(data)
+        for i, resps in enumerate(resp_all):
+            choices = resps['choices']
+            messages = resps['messages']
+            uuid = get_messages_md5(resps)
+            assert messages[-1]['role'] == 'assistant'
+            ground_truth = messages[-1]['content']
+
+            infer_requests = []
+            for decoded in choices:
+                _resps = deepcopy(resps)
+                _resps['messages'][-1]['content'] = decoded
+                infer_requests.append(_resps)
+
+            _resps = deepcopy(resps)
+            _resps['messages'][-1]['content'] = ground_truth
+            infer_requests.append(_resps)
+            if self.args.orm_model is not None:
+                orm_score, _orm_mask = self.get_orm_score(infer_requests, ground_truth)
+            else:
+                orm_score = np.array([1.0] * len(infer_requests))
+                _orm_mask = np.array([True] * len(infer_requests))
+            if self.args.prm_model is not None:
+                prm_score, _prm_mask = self.get_prm_score(infer_requests, ground_truth)
+            else:
+                prm_score = np.array([1.0] * len(infer_requests))
+                _prm_mask = np.array([True] * len(infer_requests))
+
+            _mask = _orm_mask & _prm_mask
+            if not any(_mask):
+                continue
+
+            choices.append(ground_truth)
+            choices = np.array(choices)
+
+            if self.args.orm_model is None and self.args.prm_model is None:
+                positives = choices[:-1]
+                for positive in positives:
+                    _resps = deepcopy(resps)
+                    _resps.pop('choices', None)
+                    _resps['id'] = uuid
+                    _resps['messages'][-1]['content'] = str(positive)
+                    generated.append(json.dumps(_resps, ensure_ascii=False) + '\n')
+            else:
+                score = np.array(prm_score) + np.array(orm_score * 10)
+                sorted_indices = np.argsort(score)[::-1]
+                pos_indexes = sorted_indices[0:self.args.n_best_to_keep]
+                neg_index = sorted_indices[-1]
+                pos_indexes = [int(i) for i in pos_indexes if _mask[i] and i != neg_index]
+                logger.info(
+                    f'orm:{orm_score}, prm:{prm_score}, positive index: {pos_indexes}, negative index: {neg_index}')
+                if self.args.easy_query_threshold is not None and sum([score > 0 for score in orm_score]) - 1 >= int(
+                        self.args.num_return_sequences * self.args.easy_query_threshold):
+                    continue
+                if len(pos_indexes) > 0:
+                    positives = choices[pos_indexes]
+                    negative = choices[neg_index]
+                    for positive in positives:
+                        _resps = deepcopy(resps)
+                        messages = deepcopy(messages)
+                        _resps.pop('choices', None)
+                        _resps['messages'][-1]['content'] = str(positive)
+                        _resps['rejected_response'] = str(negative)
+                        _resps['id'] = uuid
+                        generated.append(json.dumps(_resps, ensure_ascii=False) + '\n')
+        return generated