chore: 增加 h5_cache 脚本

2026-03-21 21:33:38 +08:00
parent 909f1b400c
commit 07ef471aa6
1 changed files with 86 additions and 0 deletions
@@ -0,0 +1,86 @@
+"""JSONL → H5 缓存脚本
+
+将 dataset/ 下的 JSONL 文件 tokenize 并打包为 HDF5 格式。
+
+用法:
+    python scripts/cache_h5.py pt ./dataset/chinese-c4-pretrain
+    python scripts/cache_h5.py sft ./dataset/belle-sft --pack-size 4096
+    python scripts/cache_h5.py sft ./dataset/Ling-Coder-sft --tokenizer ./my_tokenizer.json
+"""
+import argparse
+import os
+
+from pipeline import BpeTokenizer, ProcessorFactory, cache_jsonl, IOHandler
+
+
+def collect_jsonl_files(input_dir: str) -> list[str]:
+    """收集目录下所有 JSONL 文件"""
+    files = [
+        os.path.join(root, f)
+        for root, _, filenames in os.walk(input_dir)
+        for f in filenames
+        if f.endswith(".jsonl")
+    ]
+    files.sort()
+    return files
+
+
+def main():
+    parser = argparse.ArgumentParser(description="JSONL → H5 缓存")
+    parser.add_argument("type", choices=["pt", "sft", "dpo"], help="处理器类型")
+    parser.add_argument("input_dir", help="JSONL 文件所在目录")
+    parser.add_argument("-o", "--output-dir", default=None,
+                        help="H5 输出目录 (默认: <input_dir>/cached)")
+    parser.add_argument("-t", "--tokenizer", default="./tokenizer.json",
+                        help="Tokenizer 路径 (默认: ./tokenizer.json)")
+    parser.add_argument("-p", "--pack-size", type=int, default=-1,
+                        help="序列打包长度，<=0 不打包 (默认: -1)")
+    parser.add_argument("--pad-value", type=int, default=1,
+                        help="打包填充值 (默认: 1 即 <eos>)")
+    args = parser.parse_args()
+
+    # 收集 JSONL 文件
+    jsonl_files = collect_jsonl_files(args.input_dir)
+    if not jsonl_files:
+        print(f"[ERROR] No JSONL files found in {args.input_dir}")
+        return
+
+    print(f"Found {len(jsonl_files)} JSONL files:")
+    for f in jsonl_files:
+        print(f"  - {f}")
+
+    # 加载 tokenizer
+    if not os.path.exists(args.tokenizer):
+        print(f"[ERROR] Tokenizer not found: {args.tokenizer}")
+        return
+    tokenizer = BpeTokenizer(args.tokenizer)
+    print(f"Tokenizer loaded: vocab_size={len(tokenizer)}")
+
+    # 创建处理器
+    processor = ProcessorFactory.create(args.type, tokenizer)
+    print(f"Processor: {args.type} ({processor.__class__.__name__})")
+    print(f"Output keys: {processor.output_keys}")
+
+    # 输出目录
+    output_dir = args.output_dir or os.path.join(args.input_dir, "cached")
+
+    # 执行缓存
+    print(f"\nStart caching...")
+    if args.pack_size > 0:
+        print(f"  pack_size={args.pack_size}, pad_value={args.pad_value}")
+    else:
+        print(f"  no packing")
+
+    output_files = cache_jsonl(
+        files=jsonl_files,
+        output_dir=output_dir,
+        processor=processor,
+        pack_size=args.pack_size,
+        pad_value=args.pad_value,
+    )
+
+    print(f"\nDone! {len(output_files)} H5 files saved to {output_dir}")
+
+
+if __name__ == "__main__":
+    main()