fix: cache_h5 移除重复 --batch-size 参数，默认输出格式改为 bin

2026-08-19 18:01:06 +08:00
parent d67f686f10
commit 995a015c23
2 changed files with 69 additions and 8 deletions
@@ -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()