| 50 | self.num_gpus = torch.cuda.device_count() |
| 51 | |
| 52 | def process_data(self, filepath): |
| 53 | input_ids = torch.from_numpy(np.load(os.path.join(filepath, 'inputs_pack.npy'))) |
| 54 | labels = torch.from_numpy(np.load(os.path.join(filepath, 'labels_pack.npy'))) |
| 55 | weights = torch.from_numpy(np.load(os.path.join(filepath, 'weights_pack.npy'))) |
| 56 | attention_masks = json.load(open(os.path.join(filepath, 'attention_masks_pack.json'))) |
| 57 | num_gpus = torch.cuda.device_count() |
| 58 | l = (input_ids.size(0) // num_gpus) * num_gpus |
| 59 | input_ids, labels, weights, attention_masks = input_ids[:l, :], labels[:l, :], weights[:l, :], attention_masks[:l] |
| 60 | nums = [weights[i*num_gpus:(i+1)*num_gpus, :].sum() for i in range(l//num_gpus)] |
| 61 | return input_ids, attention_masks, labels, weights, nums |
| 62 | |
| 63 | def __getitem__(self, idx): |
| 64 | return { |