diff --git a/scripts/cache_h5.py b/scripts/cache_h5.py index 6df5bf9..139b86d 100644 --- a/scripts/cache_h5.py +++ b/scripts/cache_h5.py @@ -66,12 +66,6 @@ def main(): default=1_000, help="Merge every N packed chunks into one tensor, <=0 to disable (default: 1000)", ) - parser.add_argument( - "--batch-size", - type=int, - default=256, - help="Records tokenized per batch (default: 256)", - ) parser.add_argument( "--log-level", default="INFO", @@ -87,9 +81,9 @@ def main(): parser.add_argument( "-f", "--output-format", - default="h5", + default="bin", choices=["h5", "bin"], - help="Output format: h5 or bin (default: h5)", + help="Output format: h5 or bin (default: bin)", ) args = parser.parse_args() diff --git a/scripts/pre_train/ultra-fineweb-l3-qa-synthetic-sample-shards.py b/scripts/pre_train/ultra-fineweb-l3-qa-synthetic-sample-shards.py new file mode 100644 index 0000000..6515305 --- /dev/null +++ b/scripts/pre_train/ultra-fineweb-l3-qa-synthetic-sample-shards.py @@ -0,0 +1,67 @@ +import os +import random + +from datasets import load_dataset +from huggingface_hub import HfApi + +from pipeline import export_dataset + +REPO = "openbmb/Ultra-FineWeb-L3" +FRACTION = 0.1 +SEED = 42 +SAVE_ARROW = False + +CONFIGS = { + "Ultra-FineWeb-L3-en-QA-Synthetic": "data/ultrafineweb_en_l3/qa/", + "Ultra-FineWeb-L3-zh-QA-Synthetic": "data/ultrafineweb_zh_l3/qa/", +} + +HF_CACHE_DIR = "./cached_pt/ultra-fineweb-l3-qa-synthetic" +OUTPUT_DIR = "./dataset" + + +def process_func(input_dict: dict): + return {"text": input_dict["content"]} + + +def main(): + api = HfApi() + for config, prefix in CONFIGS.items(): + lang = "en" if "-en-" in config else "zh" + + shards = [ + f.path + for f in api.list_repo_tree( + REPO, path_in_repo=prefix, recursive=True, repo_type="dataset" + ) + if f.path.endswith(".parquet") + ] + k = max(1, int(len(shards) * FRACTION)) + selected = random.Random(SEED).sample(shards, k) + print(f"[{config}] total shards={len(shards)}, selected={k}", flush=True) + + dataset = load_dataset( + REPO, + data_files=selected, + split="train", + cache_dir=HF_CACHE_DIR, + ) + print(f"[{config}] loaded {len(dataset)} rows", flush=True) + + if SAVE_ARROW: + arrow_dir = os.path.join( + HF_CACHE_DIR, f"arrow-{lang}" + ) + dataset.save_to_disk(arrow_dir) + print(f"[{config}] cached arrow to {arrow_dir}", flush=True) + + export_dataset( + dataset=dataset, + output_dir=OUTPUT_DIR, + output_prefix=f"ultra-fineweb-l3-{lang}-qa-synthetic-10pct-pretrain", + process_func=process_func, + ) + + +if __name__ == "__main__": + main()