diff --git a/scripts/cache_h5.py b/scripts/cache_h5.py new file mode 100644 index 0000000..e1d6d1b --- /dev/null +++ b/scripts/cache_h5.py @@ -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 输出目录 (默认: /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 即 )") + 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()