| 64 | |
| 65 | @dataclass |
| 66 | class DataCollatorForLMPackDataset(object): |
| 67 | |
| 68 | def __call__(self, instances): |
| 69 | input_ids, attention_masks = tuple([instance[key].unsqueeze(0) for instance in instances] for key in ["input_ids", "attention_mask"]) |
| 70 | batch_seq_num = instances[0]["labels"][2] |
| 71 | labels = ([instance["labels"][0].unsqueeze(0) for instance in instances], [instance["labels"][1].unsqueeze(0) for instance in instances]) |
| 72 | input_ids = torch.cat(input_ids, dim=0) |
| 73 | labels = (torch.cat(labels[0], dim=0), torch.cat(labels[1], dim=0)) |
| 74 | labels = (labels[0], labels[1].sum()/30) |
| 75 | max_length = input_ids.shape[1] |
| 76 | attention_mask = attention_masks[0].squeeze() |
| 77 | acc_length = max_length |
| 78 | for new_attention_mask in attention_masks[1:]: |
| 79 | new_attention_mask = new_attention_mask.squeeze() |
| 80 | attention_mask = torch.cat([attention_mask, new_attention_mask[1:]+acc_length], dim=0) |
| 81 | acc_length += max_length |
| 82 | return dict( |
| 83 | input_ids=input_ids, |
| 84 | attention_mask=attention_mask, |
| 85 | labels=labels |
| 86 | ) |
| 87 | |
| 88 | def make_supervised_data_module(data_args) -> Dict: |
| 89 | print("loading data...") |
no outgoing calls
no test coverage detected