wehub-resource-sync
79 lines
3.8 KiB
Python
79 lines
3.8 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import warnings
|
|
from datasets import Dataset as HfDataset
|
|
from typing import Any, Dict, Optional
|
|
|
|
from swift.dataset import RowPreprocessor
|
|
from swift.utils import get_dist_setting, get_logger
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
class KTOPreprocessor(RowPreprocessor):
|
|
|
|
def batched_preprocess(self, batched_row: Dict[str, Any], **kwargs) -> Dict[str, Any]:
|
|
batched_row = dict(batched_row)
|
|
messages = batched_row['messages']
|
|
batch_size = len(messages)
|
|
kl_messages = [messages[-1]] + messages[:-1]
|
|
|
|
kl_response = []
|
|
for i in range(batch_size):
|
|
kl_message = kl_messages[i][-1]
|
|
assert kl_message['role'] == 'assistant'
|
|
kl_response.append(kl_message['content'])
|
|
# The name rejected_response is just for convenience in processing.
|
|
batched_row['rejected_response'] = kl_response
|
|
|
|
return batched_row
|
|
|
|
|
|
def _get_kl_dataset(dataset: Optional[HfDataset],
|
|
total_batch_size: int,
|
|
num_proc: int,
|
|
seed: Optional[int] = None) -> Optional[HfDataset]:
|
|
# Shift one position to the right in each batch.
|
|
if dataset is None:
|
|
return
|
|
dataset = dataset.shuffle(seed)
|
|
return KTOPreprocessor()(dataset, batch_size=total_batch_size, num_proc=num_proc)
|
|
|
|
|
|
def prepare_kto_dataset(args, train_dataset, val_dataset):
|
|
if args.loss_type != 'apo_zero_unpaired':
|
|
world_size = get_dist_setting()[2]
|
|
if hasattr(args, 'global_batch_size') and args.global_batch_size is not None:
|
|
total_batch_size = args.global_batch_size
|
|
else:
|
|
total_batch_size = (world_size * args.per_device_train_batch_size * args.gradient_accumulation_steps)
|
|
if total_batch_size <= 1:
|
|
raise ValueError('Batch size is 1 (too small). KTO will not work properly because the KL term '
|
|
'will be equivalent to the implied reward.')
|
|
train_dataset = _get_kl_dataset(train_dataset, total_batch_size, args.dataset_num_proc, args.data_seed)
|
|
val_dataset = _get_kl_dataset(val_dataset, total_batch_size, args.dataset_num_proc, args.data_seed)
|
|
|
|
label = train_dataset['label']
|
|
num_desirable = max(sum(label), 1)
|
|
num_undesirable = max(len(label) - num_desirable, 1) # "label" is binary
|
|
|
|
if num_desirable != num_undesirable:
|
|
# The lower and upper bounds come from Eq. (8) of https://huggingface.co/papers/2402.01306
|
|
des_weight_lower_bound = round((num_undesirable * args.undesirable_weight / num_desirable) * 1, 2)
|
|
des_weight_upper_bound = round((num_undesirable * args.undesirable_weight / num_desirable) * 1.33, 2)
|
|
und_weight_lower_bound = round((num_desirable * args.desirable_weight / num_undesirable) / 1.33, 2)
|
|
und_weight_upper_bound = round((num_desirable * args.desirable_weight / num_undesirable) / 1, 2)
|
|
|
|
des_weight_in_range = des_weight_lower_bound <= args.desirable_weight <= des_weight_upper_bound
|
|
und_weight_in_range = und_weight_lower_bound <= args.undesirable_weight <= und_weight_upper_bound
|
|
|
|
if not (des_weight_in_range or und_weight_in_range):
|
|
logger.info(f'desirable_weight: {args.desirable_weight}, undesirable_weight: {args.undesirable_weight}')
|
|
warnings.warn(
|
|
f"""
|
|
You have different amounts of desirable/positive and undesirable/negative examples but the
|
|
weights on the desirable and undesirable losses don't seem to be in an ideal range. Based
|
|
on your data, we recommend EITHER desirable_weight in [{des_weight_lower_bound}, {des_weight_upper_bound}]
|
|
or undesirable_weight in [{und_weight_lower_bound}, {und_weight_upper_bound}] (but NOT BOTH).
|
|
See the documentation on how to optimally set these weights.""", UserWarning)
|
|
return train_dataset, val_dataset
|