Class PreferenceCollator

align_anything/datasets/janus/preference.py:194–221

class PreferenceCollator:

    def __init__(self, pad_token_id: int) -> None:
        """Initialize a collator."""
        self.pad_token_id = pad_token_id

    def __call__(self, samples: list[PreferenceSample]) -> PreferenceBatch:
        return_dict = {}
        current_device = get_current_device()

        input_ids = [sample['better_input_ids'] for sample in samples] + [
            sample['worse_input_ids'] for sample in samples
        ]  # size = (2 * B, L)
        return_dict['input_ids'] = right_padding(input_ids, padding_value=self.pad_token_id).to(
            current_device
        )  # size = (2 * B, L)

        if 'pixel_values' in samples[0].keys():
            pixel_values = [sample['pixel_values'] for sample in samples]
            return_dict['pixel_values'] = torch.cat(pixel_values + pixel_values, dim=0).to(current_device).unsqueeze(0)
        
        if "better_images_emb_mask" in samples[0].keys():
            better_images_emb_mask  = [sample['better_images_emb_mask'] for sample in samples]
            worse_images_emb_mask = [sample['worse_images_emb_mask'] for sample in samples]
            return_dict['images_emb_mask'] = right_padding(better_images_emb_mask + worse_images_emb_mask, padding_value=0).to(current_device).detach()
            
        return_dict['task'] = samples[0]['task']
        return return_dict