| 568 | |
| 569 | |
| 570 | class BlockDataset(data.Dataset): |
| 571 | def __init__(self, ds, tokenizer, |
| 572 | max_seq_len=1024, |
| 573 | sample_across_doc=True, |
| 574 | non_sentence_start=0.0, filter_english=False, **kwargs): |
| 575 | """ |
| 576 | sentence_start: the stripped article must start with a complete sentence |
| 577 | """ |
| 578 | self.ds = ds |
| 579 | self.ds_len = len(self.ds) |
| 580 | self.num_samples = 1000 * self.ds_len |
| 581 | self.max_seq_len = max_seq_len |
| 582 | self.tokenizer = tokenizer |
| 583 | self.sample_across_doc = sample_across_doc |
| 584 | self.non_sentence_start = non_sentence_start |
| 585 | self.filter_english = filter_english |
| 586 | self.weighting, self.total_len = None, None |
| 587 | self.is_lazy = False |
| 588 | if self.filter_english: |
| 589 | import fasttext |
| 590 | self.model = fasttext.load_model('/mnt/lid.176.bin') |
| 591 | print_rank_0("Load language detection model") |
| 592 | if hasattr(self.ds, 'is_lazy') and self.ds.is_lazy: |
| 593 | self.is_lazy = True |
| 594 | self.init_weighting() |
| 595 | |
| 596 | def init_weighting(self): |
| 597 | if self.is_lazy: |
| 598 | lens = np.array([self.ds.get_text_len(idx) for idx in range(len(self.ds))]) |
| 599 | else: |
| 600 | lens = np.array([len(d['text']) if isinstance(d, dict) else len(d) for d in self.ds]) |
| 601 | self.total_len = np.sum(lens) |
| 602 | print_rank_0( |
| 603 | f"Dataset document count {len(lens)}, token count {self.total_len}, non sentence start{self.non_sentence_start}") |
| 604 | self.weighting = list(accumulate(lens)) |
| 605 | |
| 606 | def get_weighted_samples(self, np_rng): |
| 607 | while True: |
| 608 | idx = np_rng.randint(self.total_len) |
| 609 | data_idx = bisect_right(self.weighting, idx) |
| 610 | tokens, loss_mask = self.getidx(data_idx) |
| 611 | if self.filter_english: |
| 612 | text = self.tokenizer.DecodeIds(tokens[:1024]) |
| 613 | lang = self.model.predict(text.replace('\n', ''))[0][0] |
| 614 | if lang == '__label__en': |
| 615 | break |
| 616 | else: |
| 617 | break |
| 618 | return tokens, loss_mask |
| 619 | |
| 620 | def __len__(self): |
| 621 | return self.num_samples |
| 622 | |
| 623 | def __getitem__(self, idx): |
| 624 | # init rng |
| 625 | rng = random.Random(idx) |
| 626 | rng = np.random.RandomState(seed=[rng.randint(0, 2 ** 32 - 1) for _ in range(16)]) |
| 627 | |