| 551 | |
| 552 | |
| 553 | class ExtractionDataset(torch.utils.data.Dataset): |
| 554 | def __init__(self, args, split, tokenizer): |
| 555 | self.args = args |
| 556 | task, data_dir = args.task.lower(), args.data_dir |
| 557 | self.max_src_length, self.max_tgt_length = args.src_seq_length, args.tgt_seq_length |
| 558 | self.split = split |
| 559 | self.tokenizer = tokenizer |
| 560 | if split == "train": |
| 561 | filename = "train" |
| 562 | elif split == "dev": |
| 563 | filename = "valid" |
| 564 | elif split == "test": |
| 565 | filename = "test" |
| 566 | else: |
| 567 | raise NotImplementedError(split) |
| 568 | print_rank_0(f"Creating {task}-{split} dataset from {data_dir}") |
| 569 | self.dataset_name = split |
| 570 | source_texts, target_texts = [], [] |
| 571 | with open(os.path.join(data_dir, f"{filename}.source"), |
| 572 | encoding='utf-8') as file: |
| 573 | for line in file: |
| 574 | line = line.strip() |
| 575 | source_texts.append(line) |
| 576 | with open(os.path.join(data_dir, f"{filename}.target"), encoding='utf-8') as file: |
| 577 | for line in file: |
| 578 | line = line.strip() |
| 579 | target_texts.append(line) |
| 580 | self.examples, self.example_list = {}, [] |
| 581 | for idx, (source_text, target_text) in enumerate(zip(source_texts, target_texts)): |
| 582 | if (idx + 1) % 20000 == 0: |
| 583 | print_rank_0(f"Complete {idx + 1} examples") |
| 584 | guid = "%s-%s" % (split, idx) |
| 585 | meta = {"ref": target_text} |
| 586 | example = InputExample(guid=guid, text_a=source_text, text_b=target_text, meta=meta) |
| 587 | self.examples[guid] = example |
| 588 | self.example_list.append(example) |
| 589 | print_rank_0(f"Return {len(self.examples)} {split} examples") |
| 590 | |
| 591 | def __len__(self): |
| 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 |
no outgoing calls
no test coverage detected