| |
| import os |
| from copy import deepcopy |
| from typing import Any, Dict, List |
|
|
| import json |
| import numpy as np |
|
|
| from swift.llm import RequestConfig |
| from swift.llm.sampling.base import Sampler |
| from swift.llm.template.template_inputs import InferRequest |
| from swift.utils import get_logger |
| from .utils import get_messages_md5, get_reward |
|
|
| logger = get_logger() |
|
|
|
|
| class VanillaSampler(Sampler): |
|
|
| def __init__(self, *args, **kwargs): |
| super().__init__(*args, **kwargs) |
|
|
| if self.args.sampler_engine == 'pt': |
| from swift.llm import PtEngine |
| _Engine = PtEngine |
| elif self.args.sampler_engine == 'vllm': |
| from swift.llm import VllmEngine |
| _Engine = VllmEngine |
| elif self.args.sampler_engine == 'lmdeploy': |
| from swift.llm 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, **self.args.engine_kwargs) |
| self.infer_engine.default_template = self.template |
| self.infer_engine.strict = False |
| self.caches = self.read_cache() |
|
|
| 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') 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']]) |
|
|
| 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 |
|
|
| 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.orm_model is not None: |
| orm_score, _orm_mask = get_reward( |
| self.orm_model, infer_requests, ground_truths=[ground_truth] * len(infer_requests), threshold=0.0) |
| else: |
| orm_score = np.array([1.0] * len(infer_requests)) |
| _orm_mask = np.array([True] * len(infer_requests)) |
| if self.prm_model is not None: |
| prm_score, _prm_mask = get_reward( |
| self.prm_model, |
| infer_requests, |
| ground_truths=[ground_truth] * len(infer_requests), |
| threshold=self.args.prm_threshold) |
| 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.orm_model is None and self.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] |
| pos_indexes = [i for i in pos_indexes if _mask[i]] |
| neg_index = sorted_indices[-1] |
| 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 |
|
|
