feat : 添加 JSONL 预处理管线

- Pipeline 模板, Reader 加 transform 加 Writer 可组合 - 自动检测 JSONL 格式, 支持 messages 文本 prompt 加 response 三种 - chat 数据通过 apply_chat_template 适配, 自动生成 loss_mask - 输出对齐 Store 和 DatasetFactory, 直接用于训练 - 默认 bin 格式, CLI 入口 scripts/tools/preprocess.py
2026-05-30 17:12:42 +08:00
parent a923e0a23a
commit 138c5bcc08
2 changed files with 381 additions and 0 deletions
@@ -0,0 +1,110 @@
+"""CLI: raw JSONL → tokenized .h5/.bin via Pipeline."""
+
+import argparse
+import sys
+
+from astrai.preprocess import Pipeline, detect_format
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Raw JSONL → tokenized .h5/.bin for training"
+    )
+    parser.add_argument(
+        "inputs", nargs="+", metavar="JSONL", help="One or more JSONL files"
+    )
+    parser.add_argument(
+        "--output_dir",
+        "-o",
+        required=True,
+        help="Output directory (domain subdirs auto-created)",
+    )
+    parser.add_argument(
+        "--tokenizer_path",
+        default="params",
+        help="Path to tokenizer (default: params)",
+    )
+    parser.add_argument(
+        "--text_key",
+        default=None,
+        help="JSON key for text (auto-detect if omitted)",
+    )
+    parser.add_argument(
+        "--domain_key",
+        default=None,
+        help="JSON key for domain label (auto-detect if omitted)",
+    )
+    parser.add_argument(
+        "--max_len",
+        type=int,
+        default=2048,
+        help="Max token length per doc (default: 2048)",
+    )
+    parser.add_argument(
+        "--min_text_len",
+        type=int,
+        default=50,
+        help="Min chars per doc (default: 50)",
+    )
+    parser.add_argument(
+        "--max_text_len",
+        type=int,
+        default=2_000_000,
+        help="Max chars per doc (default: 2000000)",
+    )
+    parser.add_argument(
+        "--no_dedup",
+        action="store_true",
+        help="Skip exact dedup",
+    )
+    parser.add_argument(
+        "--max_items",
+        type=int,
+        default=None,
+        help="Max docs to process (default: all)",
+    )
+    parser.add_argument(
+        "--max_tokens_per_shard",
+        type=int,
+        default=100_000_000,
+        help="Max tokens per .h5 shard (default: 100M)",
+    )
+    parser.add_argument(
+        "--format",
+        dest="storage_format",
+        choices=["h5", "bin"],
+        default="bin",
+        help="Output format (default: bin)",
+    )
+    parser.add_argument(
+        "--detect",
+        action="store_true",
+        help="Detect and print JSONL schema, then exit",
+    )
+    args = parser.parse_args()
+
+    if args.detect:
+        fmt = detect_format(args.inputs)
+        print(f"text key   : {fmt['text_key']}")
+        print(f"domain key : {fmt['domain_key']}")
+        print(f"chat mode  : {fmt['is_chat']}")
+        sys.exit(0)
+
+    Pipeline(
+        input_paths=args.inputs,
+        output_dir=args.output_dir,
+        tokenizer_path=args.tokenizer_path,
+        text_key=args.text_key,
+        domain_key=args.domain_key,
+        max_len=args.max_len,
+        min_text_len=args.min_text_len,
+        max_text_len=args.max_text_len,
+        dedup=not args.no_dedup,
+        max_items=args.max_items,
+        max_tokens_per_shard=args.max_tokens_per_shard,
+        storage_format=args.storage_format,
+    ).run()
+
+
+if __name__ == "__main__":
+    main()