Args: batch(`List[Dict[str, Any]]`): The input data in batch padding_to(`int`, optional): Whether padding the batch to a fixed length, if none, the batch will be padded to the `longest`
(self, batch: List[Dict[str, Any]], padding_to: Optional[int] = None)
| 893 | return torch.stack(padded_sequences) |
| 894 | |
| 895 | def data_collator(self, batch: List[Dict[str, Any]], padding_to: Optional[int] = None) -> Dict[str, Any]: |
| 896 | """ |
| 897 | Args: |
| 898 | batch(`List[Dict[str, Any]]`): The input data in batch |
| 899 | padding_to(`int`, optional): Whether padding the batch to a fixed length, if none, the batch |
| 900 | will be padded to the `longest` |
| 901 | """ |
| 902 | tokenizer = self.tokenizer |
| 903 | assert tokenizer.pad_token_id is not None |
| 904 | padding_right = self.padding_side == 'right' |
| 905 | res = {} |
| 906 | |
| 907 | if 'inputs_embeds' in batch[0]: |
| 908 | inputs_embeds = [b['inputs_embeds'] for b in batch] |
| 909 | res['inputs_embeds'] = inputs_embeds |
| 910 | res['attention_mask'] = [ |
| 911 | torch.ones((inputs_embeds[i].shape[0]), dtype=torch.int64) for i in range(len(inputs_embeds)) |
| 912 | ] |
| 913 | elif 'input_ids' in batch[0]: |
| 914 | input_ids = [torch.tensor(b['input_ids']) for b in batch] |
| 915 | res['input_ids'] = input_ids |
| 916 | res['attention_mask'] = [torch.ones(len(input_ids[i]), dtype=torch.int64) for i in range(len(input_ids))] |
| 917 | |
| 918 | for key in ['labels', 'loss_scale', 'position_ids']: |
| 919 | if key in batch[0]: |
| 920 | res[key] = [torch.tensor(b[key]) for b in batch] |
| 921 | |
| 922 | if padding_to is not None: |
| 923 | assert 'input_ids' in res |
| 924 | padding_len = padding_to - res['input_ids'][0].shape[-1] |
| 925 | if padding_len > 0: |
| 926 | for key, value in zip(['input_ids', 'attention_mask', 'labels', 'loss_scale', 'position_ids'], |
| 927 | [tokenizer.pad_token_id, 0, -100, 0., -1]): |
| 928 | if key in res: |
| 929 | res[key][0] = F.pad(res[key][0], (0, padding_len) if padding_right else (padding_len, 0), |
| 930 | 'constant', value) |
| 931 | for key, value in zip(['input_ids', 'inputs_embeds', 'attention_mask', 'labels', 'loss_scale', 'position_ids'], |
| 932 | [tokenizer.pad_token_id, 0., 0, -100, 0., -1]): |
| 933 | if key in res: |
| 934 | res[key] = self.pad_sequence(res[key], value, self.padding_side) |
| 935 | |
| 936 | if '_data' in batch[0]: |
| 937 | res['_data'] = [b['_data'] for b in batch] |
| 938 | # multimodal |
| 939 | pixel_values = [b['pixel_values'] for b in batch if b.get('pixel_values') is not None] |
| 940 | if len(pixel_values) > 0: |
| 941 | res['pixel_values'] = torch.concat(pixel_values) |
| 942 | |
| 943 | image_sizes = [b['image_sizes'] for b in batch if b.get('image_sizes') is not None] |
| 944 | if len(image_sizes) > 0: |
| 945 | res['image_sizes'] = torch.concat(image_sizes) |
| 946 | |
| 947 | pixel_values_videos = [b['pixel_values_videos'] for b in batch if b.get('pixel_values_videos') is not None] |
| 948 | if len(pixel_values_videos) > 0: |
| 949 | res['pixel_values_videos'] = torch.concat(pixel_values_videos) |
| 950 | return res |
| 951 | |
| 952 | @classmethod |
no test coverage detected