(self, args, split, tokenizer)
| 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) |
nothing calls this directly
no test coverage detected