diff --git a/pre_train/chinese-cosmopedia.py b/pre_train/chinese-cosmopedia.py index b6cd249..b0d1805 100644 --- a/pre_train/chinese-cosmopedia.py +++ b/pre_train/chinese-cosmopedia.py @@ -2,17 +2,15 @@ from datasets import load_dataset from modules.utils import process_dataset if __name__ == "__main__": - max_chunk_num = 10 chunk_size = 1000000 dataset = load_dataset( "opencsg/chinese-cosmopedia", - data_files={"train": [f"data/0000{i}.parquet" for i in range(5)]} + data_files={"train": [f"data/000{i:02d}.parquet" for i in range(25)]} ) process_dataset( dataset_dict=dataset, output_subdir="chinese-wiki-pretrain", - max_chunk_num=max_chunk_num, chunk_size=chunk_size, ) \ No newline at end of file