Spaces:
Running
Running
| import datasets | |
| import os | |
| from tqdm import tqdm | |
| import webdataset as wds | |
| import json | |
| DATASET_ROOT = "/gpfs/u/home/LMCG/LMCGljnn/scratch-shared/the_pile/all/train" | |
| OUT_DIR = "/gpfs/u/home/LMCG/LMCGljnn/scratch-shared/junyan/raw/the_pile" | |
| SAMPLE_PER_SHARD = 100000 | |
| if __name__ == "__main__": | |
| os.makedirs(OUT_DIR) | |
| print("load dataset...") | |
| pile = datasets.load_from_disk(DATASET_ROOT) | |
| total_num = pile.num_rows | |
| print("total num:", total_num) | |
| num = 0 | |
| pbar = tqdm(total=total_num) | |
| with wds.ShardWriter(OUT_DIR+"/%05d.tar", maxcount=SAMPLE_PER_SHARD, encoder=False) as sink: | |
| for sample in pile.iter(4096): | |
| for text, meta in zip(sample["text"], sample["meta"]): | |
| pbar.update(1) | |
| if meta.get("pile_set_name", None) == "Github": | |
| continue | |
| num += 1 | |
| sink.write({ | |
| '__key__': str(num), | |
| 'txt': text.encode("utf-8"), | |
| 'json': json.dumps(meta, indent=4).encode("utf-8"), | |
| }) | |
| print(f"{num} out of {total_num} is written") | |