reafactor: 重构项目

scripts/cache_h5.py

@@ -1,46 +1,36 @@
-"""JSONL → H5 缓存脚本
+"""JSONL to H5 caching script.
 
-将 dataset/ 下的 JSONL 文件 tokenize 并打包为 HDF5 格式。
+Tokenize JSONL files and pack them into HDF5 format.
 
-用法:
+Usage:
     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/belle-sft --pack-size 4096 --strategy alpaca
     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
+from pipeline import BpeTokenizer, ProcessorFactory, cache_jsonl
+from pipeline.io import IOHandler
 
 
 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 = argparse.ArgumentParser(description="JSONL -> H5 cache")
+    parser.add_argument("type", choices=["pt", "sft", "dpo"], help="Processor type")
+    parser.add_argument("input_dir", help="Directory containing JSONL files")
     parser.add_argument("-o", "--output-dir", default=None,
-                        help="H5 输出目录 (默认: <input_dir>/cached)")
+                        help="H5 output dir (default: <input_dir>/cached)")
     parser.add_argument("-t", "--tokenizer", default="./tokenizer.json",
-                        help="Tokenizer 路径 (默认: ./tokenizer.json)")
+                        help="Tokenizer path (default: ./tokenizer.json)")
+    parser.add_argument("-s", "--strategy", default=None,
+                        help="Prompt strategy: chatml, alpaca (default: chatml)")
     parser.add_argument("-p", "--pack-size", type=int, default=-1,
-                        help="序列打包长度，<=0 不打包 (默认: -1)")
+                        help="Pack size, <=0 to disable (default: -1)")
     parser.add_argument("--pad-value", type=int, default=1,
-                        help="打包填充值 (默认: 1 即 <eos>)")
+                        help="Padding value (default: 1)")
     args = parser.parse_args()
 
-    # 收集 JSONL 文件
-    jsonl_files = collect_jsonl_files(args.input_dir)
+    jsonl_files = IOHandler.fetch_files(args.input_dir, suffix=".jsonl")
     if not jsonl_files:
         print(f"[ERROR] No JSONL files found in {args.input_dir}")
         return
@@ -49,37 +39,38 @@ def main():
     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)
+    if args.strategy:
+        processor = ProcessorFactory.create_with_strategy_name(
+            args.type, tokenizer, args.strategy
+        )
+    else:
+        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(
+    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}")
+    print(f"\nDone! Output saved to {output_dir}")
 
 
 if __name__ == "__main__":