| 109 | return dict(input_ids=self.input_ids[i], labels=self.labels[i], id=i) |
| 110 | |
| 111 | def padding(inputs, padding_token, cutoff = None): |
| 112 | num_elems = len(inputs) |
| 113 | if cutoff is None: |
| 114 | cutoff = max([len(item) for item in inputs]) |
| 115 | else: |
| 116 | cutoff = min(max([len(item) for item in inputs]), cutoff) |
| 117 | |
| 118 | tokens = torch.ones(num_elems, cutoff).long().to(inputs[0].device) * padding_token |
| 119 | for i in range(num_elems): |
| 120 | toks = inputs[i] |
| 121 | length = min(cutoff, len(toks)) |
| 122 | tokens[i, -length:] = toks[-length:] |
| 123 | return tokens |
| 124 | |
| 125 | @dataclass |
| 126 | class DataCollatorForSupervisedDataset(object): |