feat: 并行 batch tokenization + cache_jsonl 批处理支持

- pipeline/tokenize/tokenizer.py: encode() 全部走 encode_batch(支持单条/批量)
- pipeline/processors/base.py: BaseProcessor 新增 process_batch()
- pipeline/processors/pretrain.py: PreTrainProcessor 覆盖 process_batch() 批量编码
- pipeline/io/export.py: cache_jsonl 新增 batch_size 参数默认 1000, 批量处理
- scripts/cache_h5.py: 新增 --batch-size 参数, 默认 tokenizer 路径改为 ../AstrAI/params

pipeline/io/export.py

@@ -116,6 +116,7 @@ def cache_jsonl(
     group_size: int = 1_000,
     pack_algo: Optional[str] = None,
     output_format: str = "h5",
+    batch_size: int = 1000,
 ) -> List[str]:
     """Tokenize JSONL files and save as HDF5 or binary.
 
@@ -133,6 +134,8 @@ def cache_jsonl(
         pack_algo: Packing algorithm: 'bfd' (default), 'ffd',
             'greedy'. Only used when pack_size > 0.
         output_format: ``"h5"`` or ``"bin"``.
+        batch_size: Number of lines to batch-process together for parallel
+            tokenization via encode_batch (default: 1000).
 
     Returns:
         List of generated file paths.
@@ -157,34 +160,45 @@ def cache_jsonl(
         arrows_batch: Dict[str, List] = {key: [] for key in output_keys}
         batch_tokens: int = 0
 
+        buf: List[str] = []
+        buf_num: int = 0
+
+        def flush_buf():
+            nonlocal batch_tokens
+            if not buf:
+                return
+            samples = []
+            for line in buf:
+                try:
+                    samples.append(json.loads(line))
+                except json.JSONDecodeError as e:
+                    logger.warning(f"JSON decode error, skipping: {e}")
+            buf.clear()
+            if not samples:
+                return
+            results = processor.process_batch(samples) if hasattr(processor, "process_batch") else [processor.process(s) for s in samples]
+            for result in results:
+                if result is not None:
+                    for key in output_keys:
+                        arrows_batch[key].append(result[key])
+                    if target_tokens > 0:
+                        batch_tokens += int(result[output_keys[0]].shape[0])
+
         with open(file_path, "r", encoding="utf-8") as f:
             for line_num, line in enumerate(
                 tqdm(f, desc=f"Processing {file_name}", leave=False), start=1
             ):
-                try:
-                    result = processor.process(json.loads(line))
-                    if result is not None:
+                buf.append(line)
+                if len(buf) >= batch_size:
+                    flush_buf()
+                    if target_tokens > 0 and batch_tokens >= target_tokens:
+                        packed = pack_tensors(arrows_batch, pack_size, pad_value, dtypes, pad_values=pad_values, algo=pack_algo)
                         for key in output_keys:
-                            arrows_batch[key].append(result[key])
-                        if target_tokens > 0:
-                            batch_tokens += int(result[output_keys[0]].shape[0])
-                except json.JSONDecodeError as e:
-                    logger.warning(
-                        f"JSON decode error in {file_path} line {line_num}: {e}. Skipping line."
-                    )
-                    continue
-                except Exception as e:
-                    logger.warning(
-                        f"Unexpected error processing line {line_num} in {file_path}: {e}. Skipping line."
-                    )
-                    continue
+                            all_packed[key].extend(packed[key])
+                            arrows_batch[key] = []
+                        batch_tokens = 0
 
-                if target_tokens > 0 and batch_tokens >= target_tokens:
-                    packed = pack_tensors(arrows_batch, pack_size, pad_value, dtypes, pad_values=pad_values, algo=pack_algo)
-                    for key in output_keys:
-                        all_packed[key].extend(packed[key])
-                        arrows_batch[key] = []
-                    batch_tokens = 0
+        flush_buf()
 
         if arrows_batch[output_keys[0]]:
             if pack_size > 0:

pipeline/processors/base.py

@@ -82,6 +82,21 @@ class BaseProcessor(ABC):
         """Return list of output tensor key names."""
         pass
 
+    def process_batch(self, input_dicts: List[Dict[str, Any]]) -> List[Dict[str, Tensor]]:
+        """Process a batch of input samples.
+
+        Default implementation calls process() for each sample.
+        Subclasses should override for efficient batch processing
+        (e.g., using tokenizer.encode_batch).
+
+        Args:
+            input_dicts: List of input dictionaries.
+
+        Returns:
+            List of output dictionaries mapping output key names to tensors.
+        """
+        return [self.process(d) for d in input_dicts]
+
     def validate_input(self, input_dict: Dict[str, Any]) -> None:
         """Validate input against schema before processing.
 

pipeline/processors/pretrain.py

@@ -43,6 +43,14 @@ class PreTrainProcessor(BaseProcessor):
         tokens = self.tokenizer.encode(f"{segment}{self._eos_token}")
         return {"sequence": torch.tensor(tokens, dtype=torch.int32)}
 
+    def process_batch(self, input_dicts: List[Dict[str, Any]]) -> List[Dict[str, Tensor]]:
+        texts = [f"{d['text']}{self._eos_token}" for d in input_dicts]
+        batch_tokens = self.tokenizer.encode(texts)
+        return [
+            {"sequence": torch.tensor(tokens, dtype=torch.int32)}
+            for tokens in batch_tokens
+        ]
+
     @property
     def output_keys(self) -> List[str]:
         return ["sequence"]

pipeline/tokenize/tokenizer.py

@@ -240,22 +240,18 @@ class AutoTokenizer:
                 "Tokenizer not initialized. Load or create a tokenizer first."
             )
 
-        if isinstance(tokens, str):
-            encoded = self._tokenizer.encode(
-                tokens,
-                is_pretokenized=is_pretokenized,
-                add_special_tokens=add_special_tokens,
-            )
-            return encoded.ids if out_ids else encoded.tokens
-        else:
-            encoded_list = self._tokenizer.encode_batch(
-                tokens,
-                is_pretokenized=is_pretokenized,
-                add_special_tokens=add_special_tokens,
-            )
-            return [
-                encoded.ids if out_ids else encoded.tokens for encoded in encoded_list
-            ]
+        single = isinstance(tokens, str)
+        if single:
+            tokens = [tokens]
+        encoded_list = self._tokenizer.encode_batch(
+            tokens,
+            is_pretokenized=is_pretokenized,
+            add_special_tokens=add_special_tokens,
+        )
+        result = [
+            encoded.ids if out_ids else encoded.tokens for encoded in encoded_list
+        ]
+        return result[0] if single else result
 
     def decode(self, tokens: List[int], skip_special_tokens: bool = True) -> str:
         """Decode token IDs to text."""

scripts/cache_h5.py

@@ -72,6 +72,12 @@ def main():
         choices=["DEBUG", "INFO", "WARNING", "ERROR"],
         help="Logging level (default: INFO)",
     )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=1000,
+        help="Lines per batch for parallel tokenization via encode_batch (default: 1000)",
+    )
     parser.add_argument(
         "-f",
         "--output-format",
@@ -133,6 +139,7 @@ def main():
         group_size=args.group_size,
         pack_algo=args.pack_algo,
         output_format=args.output_format,
+        batch_size=args.batch_size,
     )
     print(f"\nDone! Output saved to {output_dir}")