Collate examples for supervised fine-tuning.
| 167 | |
| 168 | @dataclass |
| 169 | class DataCollatorForSupervisedDataset(object): |
| 170 | """Collate examples for supervised fine-tuning.""" |
| 171 | |
| 172 | tokenizer: transformers.PreTrainedTokenizer |
| 173 | |
| 174 | def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]: |
| 175 | input_ids, labels, ids = tuple([instance[key] for instance in instances] for key in ("input_ids", "labels", 'id')) |
| 176 | input_ids = padding(input_ids, self.tokenizer.pad_token_id, cutoff = 256) |
| 177 | labels = padding(labels, IGNORE_INDEX, cutoff = 256) |
| 178 | |
| 179 | return dict( |
| 180 | input_ids=input_ids, |
| 181 | labels=labels, |
| 182 | id=torch.tensor(ids).to(input_ids.device), |
| 183 | attention_mask=input_ids.ne(self.tokenizer.pad_token_id), |
| 184 | ) |
| 185 | |
| 186 | |
| 187 | def make_supervised_data_module(tokenizer: transformers.PreTrainedTokenizer, data_path) -> Dict: |
no outgoing calls