collator.py

from dataclasses import dataclass
from typing import Any, Dict, Sequence

import torch
from transformers import DataCollatorForSeq2Seq


@dataclass
class PairwiseDataCollatorWithPadding(DataCollatorForSeq2Seq):
    r"""
    Data collator for pairwise data.
    """

    def __call__(
            self, features: Sequence[Dict[str,
                                          Any]]) -> Dict[str, 'torch.Tensor']:
        r"""
        Pads batched data to the longest sequence in the batch.

        We generate 2 * n examples where the first n examples represent chosen examples and
        the last n examples represent rejected examples.
        """
        concatenated_features = []
        for key in ('chosen', 'rejected'):
            for feature in features:
                target_feature = {
                    'input_ids': feature['{}_input_ids'.format(key)],
                    'attention_mask': feature['{}_attention_mask'.format(key)],
                    'labels': feature['{}_labels'.format(key)],
                }
                if 'pixel_values' in feature:
                    target_feature['pixel_values'] = feature['pixel_values']

                if '{}_token_type_ids'.format(key) in feature:
                    target_feature['token_type_ids'] = feature[
                        '{}_token_type_ids'.format(key)]

                concatenated_features.append(target_feature)

        return super().__call__(concatenated_features)


@dataclass
class KTODataCollatorWithPadding(DataCollatorForSeq2Seq):
    r"""
    Data collator for KTO data.
    """

    def __call__(
            self, features: Sequence[Dict[str,
                                          Any]]) -> Dict[str, 'torch.Tensor']:
        target_features = []
        kl_features = []
        kto_tags = []
        for feature in features:
            target_feature = {
                'input_ids': feature['input_ids'],
                'attention_mask': feature['attention_mask'],
                'labels': feature['labels'],
            }
            kl_feature = {
                'input_ids': feature['kl_input_ids'],
                'attention_mask': feature['kl_attention_mask'],
                'labels': feature['kl_labels'],
            }
            if 'pixel_values' in feature:
                target_feature['pixel_values'] = feature['pixel_values']

            if 'token_type_ids' in feature:
                target_feature['token_type_ids'] = feature['token_type_ids']
                kl_feature['token_type_ids'] = feature['kl_token_type_ids']

            target_features.append(target_feature)
            kl_features.append(kl_feature)
            kto_tags.append(feature['kto_tags'])

        batch = super().__call__(target_features)
        kl_batch = super().__call__(kl_features)
        batch['kl_input_ids'] = kl_batch['input_ids']
        batch['kl_attention_mask'] = kl_batch['attention_mask']
        batch['kl_labels'] = kl_batch['labels']
        if 'token_type_ids' in batch:
            batch['kl_token_type_ids'] = kl_batch['token_type_ids']

        batch['kto_tags'] = torch.tensor(kto_tags)
        return batch