| 97 | Make dataset and collator for supervised fine-tuning or continue pre-train. |
| 98 | """ |
| 99 | def load_data(dataset_name): |
| 100 | if dataset_name == 'alpaca': |
| 101 | return load_dataset("tatsu-lab/alpaca") |
| 102 | elif dataset_name == 'oasst1': |
| 103 | return load_dataset("timdettmers/openassistant-guanaco") |
| 104 | elif dataset_name == 'deita-6k': |
| 105 | dataset = load_dataset("hkust-nlp/deita-6k-v0", split = "train") |
| 106 | dataset = [row for row in dataset] |
| 107 | return dataset |
| 108 | elif dataset_name == 'deita-10k': |
| 109 | dataset = load_dataset("hkust-nlp/deita-10k-v0", split = "train") |
| 110 | dataset = [row for row in dataset] |
| 111 | return dataset |
| 112 | elif dataset_name == 'c4': |
| 113 | try: |
| 114 | # load from local file, a fast manner |
| 115 | dataset = load_dataset("arrow", |
| 116 | data_files={ |
| 117 | "train": "/cpfs01/user/chenmengzhao/huggingface/datasets/allenai___json/allenai--c4-6fbe877195f42de5/0.0.0/0f7e3662623656454fcd2b650f34e886a7db4b9104504885bd462096cc7a9f51/json-train-00000-of-00002.arrow", |
| 118 | "validation": "/cpfs01/user/chenmengzhao/huggingface/datasets/allenai___json/allenai--c4-efc3d4f4606f44bd/0.0.0/fe5dd6ea2639a6df622901539cb550cf8797e5a6b2dd7af1cf934bed8e233e6e/json-validation.arrow", |
| 119 | }, |
| 120 | ) |
| 121 | except: |
| 122 | dataset = load_dataset("allenai/c4","allenai--c4", |
| 123 | data_files={ |
| 124 | "train": "en/c4-train.00000-of-01024.json.gz", |
| 125 | "validation": "en/c4-validation.00000-of-00008.json.gz", |
| 126 | }, |
| 127 | ) |
| 128 | return dataset |
| 129 | elif dataset_name == 'redpajama': |
| 130 | try: |
| 131 | loacal_dataset = "/cpfs01/user/chenmengzhao/huggingface/datasets/togethercomputer___red_pajama-data-1_t-sample" |
| 132 | dataset = load_dataset(loacal_dataset) |
| 133 | except: |
| 134 | dataset = load_dataset("togethercomputer/RedPajama-Data-1T-Sample") |
| 135 | if "validation" not in dataset.keys(): |
| 136 | validation_split = args.eval_dataset_size |
| 137 | dataset["validation"] = load_dataset( |
| 138 | loacal_dataset, |
| 139 | split=f"train[:{validation_split}]", |
| 140 | ) |
| 141 | dataset["train"] = load_dataset( |
| 142 | loacal_dataset, |
| 143 | split=f"train[{validation_split}:]", |
| 144 | ) |
| 145 | return dataset |
| 146 | else: |
| 147 | raise NotImplementedError(f"Dataset {dataset_name} not implemented yet.") |
| 148 | |
| 149 | |
| 150 | def format_dataset(dataset, dataset_format): |