(self, idx)
| 592 | return len(self.example_list) |
| 593 | |
| 594 | def __getitem__(self, idx): |
| 595 | example = self.example_list[idx] |
| 596 | source_text, target_text = example.text_a, example.text_b |
| 597 | mask_token = 'MASK' |
| 598 | mask_id = self.tokenizer.get_command(mask_token).Id |
| 599 | sop_id = self.tokenizer.get_command('sop').Id |
| 600 | eop_id = self.tokenizer.get_command('eop').Id |
| 601 | pad_id = self.tokenizer.get_command('pad').Id |
| 602 | |
| 603 | def pad_to(text, max_len, pad_id): |
| 604 | if len(text) > max_len: |
| 605 | text = text[:max_len] |
| 606 | else: |
| 607 | text = text + [pad_id] * (max_len - len(text)) |
| 608 | return text |
| 609 | |
| 610 | source_tokens = self.tokenizer.EncodeAsIds(source_text).tokenization |
| 611 | masked_tgt = target_text.split("|") |
| 612 | source_tokens = pad_to(source_tokens, self.max_src_length, pad_id) |
| 613 | sep = len(source_tokens) |
| 614 | position_ids = list(range(len(source_tokens))) |
| 615 | block_position_ids = [0] * len(source_tokens) |
| 616 | if self.split == 'train': |
| 617 | mask_positions = [i for i, x in enumerate(source_tokens) if x == mask_id] |
| 618 | assert len(mask_positions) <= len(masked_tgt) |
| 619 | tokens = source_tokens |
| 620 | target_ids = [0] * len(source_tokens) |
| 621 | loss_mask = [0] * len(source_tokens) |
| 622 | for i, mask_pos in enumerate(mask_positions): |
| 623 | tgt_text = masked_tgt[i] |
| 624 | tgt_tokens = self.tokenizer.EncodeAsIds(" " + tgt_text).tokenization |
| 625 | tokens += [sop_id] + tgt_tokens |
| 626 | target_ids += tgt_tokens + [eop_id] |
| 627 | loss_mask += [1] * (len(tgt_tokens) + 1) |
| 628 | position_ids += [mask_pos] * (len(tgt_tokens) + 1) |
| 629 | block_position_ids += [i + 1 for i in range(len(tgt_tokens) + 1)] |
| 630 | tokens = pad_to(tokens, self.max_src_length + self.max_tgt_length, pad_id) |
| 631 | target_ids = pad_to(target_ids, self.max_src_length + self.max_tgt_length, pad_id) |
| 632 | loss_mask = pad_to(loss_mask, self.max_src_length + self.max_tgt_length, 0) |
| 633 | position_ids = pad_to(position_ids, self.max_src_length + self.max_tgt_length, 0) |
| 634 | block_position_ids = pad_to(block_position_ids, self.max_src_length + self.max_tgt_length, 0) |
| 635 | position_ids = [position_ids, block_position_ids] |
| 636 | sample = {'text': np.array(tokens, dtype=np.int64), 'target': np.array(target_ids, dtype=np.int64), |
| 637 | 'attention_mask': np.array(sep, dtype=np.int64), |
| 638 | 'loss_mask': np.array(loss_mask, dtype=np.int64), |
| 639 | "position_id": np.array(position_ids, dtype=np.int64), "uid": example.guid} |
| 640 | else: |
| 641 | tokens = source_tokens + [sop_id] |
| 642 | mask_pos = source_tokens.index(mask_id) |
| 643 | position_ids = position_ids + [mask_pos] |
| 644 | block_position_ids = block_position_ids + [1] |
| 645 | position_ids = [position_ids, block_position_ids] |
| 646 | sample = {'text': np.array(tokens, dtype=np.int64), 'attention_mask': np.array(sep, dtype=np.int64), |
| 647 | "position_id": np.array(position_ids, dtype=np.int64), "uid": example.guid} |
| 648 | return sample |
| 649 | |
| 650 | |
| 651 | class BlankLMDataset(torch.utils.data.Dataset): |
nothing calls this directly
no test coverage detected