# 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