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

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()
 

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()