(dataset_name, max_sample=None, tokenizer=None)
| 59 | |
| 60 | |
| 61 | def get_gen_dataset(dataset_name, max_sample=None, tokenizer=None): |
| 62 | if dataset_name == "wikitext": |
| 63 | return get_wiki_dataset(max_sample) |
| 64 | elif dataset_name == "repajama": |
| 65 | return get_redpajama_dataset(max_sample) |
| 66 | elif dataset_name == "alpaca": |
| 67 | return get_alpaca_dataset(max_sample, tokenizer) |
| 68 | elif dataset_name == "alpaca-solar": |
| 69 | return get_alpaca_solar_dataset(max_sample, tokenizer) |
| 70 | elif dataset_name == "openorca-solar": |
| 71 | return get_openorca_solar_dataset(max_sample, tokenizer) |
| 72 | elif dataset_name == "ultra-solar": |
| 73 | return get_ultra_solar_dataset(max_sample, tokenizer) |
| 74 | elif dataset_name == "code": |
| 75 | return get_code_dataset(max_sample, tokenizer) |
| 76 | elif dataset_name == "math": |
| 77 | return get_math_dataset(max_sample, tokenizer) |
| 78 | else: |
| 79 | raise ValueError(f"{dataset_name} not implement yet") |
| 80 | |
| 81 | def extract_random_dataset(sources, targets, max_sample=None): |
| 82 | if max_sample is not None: |
no test coverage detected