perf: batch tokenizer preprocessing

2026-07-23 18:42:19 +08:00
parent 021e6f3788
commit 99b5d2b2da
5 changed files with 320 additions and 25 deletions
@@ -22,9 +22,19 @@ def main():
        default="params",
        help="Path to tokenizer directory (default: params)",
    )
+    parser.add_argument(
+        "--batch_size",
+        type=int,
+        default=None,
+        help="Number of records tokenized together (default: config value)",
+    )
    args = parser.parse_args()

    config = PipelineConfig.from_file(args.config)
+    if args.batch_size is not None:
+        if args.batch_size < 1:
+            parser.error("--batch_size must be at least 1")
+        config.preprocessing.batch_size = args.batch_size

    Pipeline(
        config=config,