Dataset containing sentencepairs for BERT training. Each index corresponds to a randomly generated sentence pair. Arguments: ds (Dataset or array-like): data corpus to use for training max_seq_len (int): maximum sequence length to use for a sentence pair mask_lm_prob
| 838 | |
| 839 | |
| 840 | class BertSentencepairDataset(data.Dataset): |
| 841 | """ |
| 842 | Dataset containing sentencepairs for BERT training. Each index corresponds to a randomly generated sentence pair. |
| 843 | Arguments: |
| 844 | ds (Dataset or array-like): data corpus to use for training |
| 845 | max_seq_len (int): maximum sequence length to use for a sentence pair |
| 846 | mask_lm_prob (float): proportion of tokens to mask for masked LM |
| 847 | max_preds_per_seq (int): Maximum number of masked tokens per sentence pair. Default: math.ceil(max_seq_len*mask_lm_prob/10)*10 |
| 848 | short_seq_prob (float): Proportion of sentence pairs purposefully shorter than max_seq_len |
| 849 | dataset_size (int): number of random sentencepairs in the dataset. Default: len(ds)*(len(ds)-1) |
| 850 | |
| 851 | """ |
| 852 | |
| 853 | def __init__(self, ds, max_seq_len=512, mask_lm_prob=.15, max_preds_per_seq=None, short_seq_prob=.01, |
| 854 | dataset_size=None, presplit_sentences=False, weighted=True, **kwargs): |
| 855 | self.ds = ds |
| 856 | self.ds_len = len(self.ds) |
| 857 | self.tokenizer = self.ds.GetTokenizer() |
| 858 | self.vocab_words = list(self.tokenizer.text_token_vocab.values()) |
| 859 | self.ds.SetTokenizer(None) |
| 860 | self.max_seq_len = max_seq_len |
| 861 | self.mask_lm_prob = mask_lm_prob |
| 862 | if max_preds_per_seq is None: |
| 863 | max_preds_per_seq = math.ceil(max_seq_len * mask_lm_prob / 10) * 10 |
| 864 | self.max_preds_per_seq = max_preds_per_seq |
| 865 | self.short_seq_prob = short_seq_prob |
| 866 | self.dataset_size = dataset_size |
| 867 | if self.dataset_size is None: |
| 868 | self.dataset_size = self.ds_len * (self.ds_len - 1) |
| 869 | self.presplit_sentences = presplit_sentences |
| 870 | if not self.presplit_sentences: |
| 871 | nltk.download('punkt', download_dir="./nltk") |
| 872 | self.weighted = weighted |
| 873 | self.get_weighting() |
| 874 | |
| 875 | def get_weighting(self): |
| 876 | if self.weighted: |
| 877 | if hasattr(self.ds, 'is_lazy') and self.ds.is_lazy: |
| 878 | lens = np.array(self.ds.lens) |
| 879 | else: |
| 880 | lens = np.array([len(d['text']) if isinstance(d, dict) else len(d) for d in self.ds]) |
| 881 | self.total_len = np.sum(lens) |
| 882 | self.weighting = list(accumulate(lens)) |
| 883 | else: |
| 884 | self.weighting = None |
| 885 | |
| 886 | def get_weighted_samples(self, np_rng): |
| 887 | if self.weighting is not None: |
| 888 | idx = np_rng.randint(self.total_len) |
| 889 | return bisect_right(self.weighting, idx) |
| 890 | else: |
| 891 | return np_rng.randint(self.ds_len) |
| 892 | |
| 893 | def __len__(self): |
| 894 | return self.dataset_size |
| 895 | |
| 896 | def __getitem__(self, idx): |
| 897 | # get rng state corresponding to index (allows deterministic random pair) |