(path, process_fn, columns=None, cache_dir='~/.cache/huggingface/datasets', offline=False, transformer_name = None, rebuild=False)
| 19 | return first_name, second_name, split |
| 20 | |
| 21 | def load_hf_dataset(path, process_fn, columns=None, cache_dir='~/.cache/huggingface/datasets', offline=False, transformer_name = None, rebuild=False): |
| 22 | dataset_name, sub_name, split = parse_huggingface_path(path) |
| 23 | datasets.config.HF_DATASETS_OFFLINE = int(offline) |
| 24 | if transformer_name: |
| 25 | dataset_path = cache_dir + '/' + dataset_name + "_" + sub_name + "_" + split + "_" + transformer_name + ".data" |
| 26 | else: |
| 27 | dataset_path = None |
| 28 | |
| 29 | if dataset_path and os.path.exists(dataset_path) and not rebuild: |
| 30 | dataset = datasets.load_from_disk(dataset_path) |
| 31 | else: |
| 32 | dataset = load_dataset(dataset_name, sub_name, cache_dir=cache_dir, split=split, |
| 33 | download_config=datasets.utils.DownloadConfig(max_retries=20)) # TODO |
| 34 | # dataset = dataset.filter(lambda example, indice: indice % 100 == 0, with_indices=True) |
| 35 | print_rank0(f'> Preprocessing the {dataset_name} by process_fn... Next time will return cached files.\n> Pass "rebuild=True" to load_hf_dataset if change process_fn. Change "transformer_name" for different tokenizers or models.') |
| 36 | dataset = dataset.map(process_fn, batched=False, load_from_cache_file=True) |
| 37 | if dataset_path: |
| 38 | dataset.save_to_disk(dataset_path) |
| 39 | dataset.set_format(type='torch', columns=columns) |
| 40 | return dataset |
no test coverage detected