File size: 3,679 Bytes
7feac49 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
# Copyright (c) Alibaba, Inc. and its affiliates.
import warnings
from typing import Any, Dict, Optional
from datasets import Dataset as HfDataset
from swift.utils import get_dist_setting, get_logger
from ..dataset import RowPreprocessor
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):
world_size = get_dist_setting()[2]
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
|