wehub-resource-sync
234 lines
9.5 KiB
Python
234 lines
9.5 KiB
Python
# 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
|