| 713 | |
| 714 | |
| 715 | class GPT2Dataset(data.Dataset): |
| 716 | |
| 717 | def __init__(self, ds, tokenizer, |
| 718 | max_seq_len=1024, |
| 719 | num_samples=None, |
| 720 | weighted=True, |
| 721 | sample_across_doc=True, |
| 722 | random_across_doc_sampling=True, |
| 723 | sentence_start=False, **kwargs): |
| 724 | """ |
| 725 | sentence_start: the stripped article must start with a complete sentence |
| 726 | """ |
| 727 | self.ds = ds |
| 728 | self.ds_len = len(self.ds) |
| 729 | self.num_samples = num_samples |
| 730 | if num_samples is None: |
| 731 | self.num_samples = 1000 * self.ds_len |
| 732 | self.max_seq_len = max_seq_len |
| 733 | self.tokenizer = tokenizer |
| 734 | self.weighted = weighted |
| 735 | self.sample_across_doc = sample_across_doc |
| 736 | self.random_across_doc_sampling = random_across_doc_sampling |
| 737 | self.sentence_start = sentence_start |
| 738 | self.weighting, self.total_len = None, None |
| 739 | self.is_lazy = False |
| 740 | if hasattr(self.ds, 'is_lazy') and self.ds.is_lazy: |
| 741 | self.is_lazy = True |
| 742 | self.init_weighting() |
| 743 | |
| 744 | def init_weighting(self): |
| 745 | if self.weighted: |
| 746 | if self.is_lazy: |
| 747 | lens = np.array([self.ds.get_text_len(idx) for idx in range(len(self.ds))]) |
| 748 | else: |
| 749 | lens = np.array([len(d['text']) if isinstance(d, dict) |
| 750 | else len(d) for d in self.ds]) |
| 751 | self.total_len = np.sum(lens) |
| 752 | print_rank_0(f"Dataset document count {len(lens)}, token count {self.total_len}") |
| 753 | self.weighting = list(accumulate(lens)) |
| 754 | else: |
| 755 | self.weighting = None |
| 756 | |
| 757 | def get_weighted_samples(self, np_rng): |
| 758 | if self.weighting is not None: |
| 759 | idx = np_rng.randint(self.total_len) |
| 760 | return bisect_right(self.weighting, idx) |
| 761 | else: |
| 762 | return np_rng.randint(self.ds_len) |
| 763 | |
| 764 | def __len__(self): |
| 765 | return self.num_samples |
| 766 | |
| 767 | def __getitem__(self, idx): |
| 768 | # init rng |
| 769 | rng = random.Random(idx) |
| 770 | rng = np.random.RandomState(seed=[rng.randint(0, 2 ** 32 - 1) for _ in range(16)]) |
| 771 | |
| 772 | # get possibly weighted random index from dataset |