(dataset)
| 182 | |
| 183 | # Split dataset into train/val/test (and wrap bert dataset) |
| 184 | def wrap_dataset(dataset): |
| 185 | if ds_type.lower() == 'bert': |
| 186 | presplit_sentences = kwargs['presplit_sentences'] if 'presplit_sentences' in kwargs else False |
| 187 | dataset = BertSentencepairDataset(dataset, max_seq_len=seq_length, presplit_sentences=presplit_sentences) |
| 188 | elif ds_type.lower() == 'gpt-xl': |
| 189 | assert pre_tokenize |
| 190 | dataset = XLDataset(dataset, tokenizer, max_seq_len=seq_length, mem_len=mem_length, |
| 191 | sample_across_doc=not sample_one_document) |
| 192 | elif ds_type.lower() == 'gpt2': |
| 193 | dataset = GPT2Dataset(dataset, tokenizer, max_seq_len=seq_length, sample_across_doc=not sample_one_document) |
| 194 | elif ds_type.lower() == 'block': |
| 195 | dataset = BlockDataset(dataset, tokenizer, max_seq_len=seq_length, |
| 196 | sample_across_doc=not sample_one_document, filter_english=filter_english, |
| 197 | non_sentence_start=non_sentence_start) |
| 198 | return dataset |
| 199 | |
| 200 | if should_split(split): |
| 201 | ds = split_ds(ds, split, shuffle=shuffle, save_splits=save_splits, load_splits=load_splits) |
no test coverage detected