| 54 | |
| 55 | @staticmethod |
| 56 | def pet_wrapper(data): |
| 57 | text = data['text'] |
| 58 | loss_mask = data['logit_mask'] |
| 59 | target = data['target'] |
| 60 | attention_mask = data['mask'] |
| 61 | position_id = data['position'] |
| 62 | label = data['label'] |
| 63 | if len(text.shape) == 2: |
| 64 | text = text[label] |
| 65 | loss_mask = loss_mask[label] |
| 66 | target = target[label] |
| 67 | attention_mask = attention_mask[label] |
| 68 | position_id = position_id[label] |
| 69 | else: |
| 70 | target = target[label] |
| 71 | if not target.shape: |
| 72 | target = target.repeat(len(text)) |
| 73 | return {'text': text, 'target': target, 'loss_mask': loss_mask, 'position_id': position_id, |
| 74 | 'attention_mask': attention_mask} |
| 75 | |
| 76 | def __getitem__(self, idx): |
| 77 | if self.reweight: |