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
This commit is contained in:
2026-07-25 12:19:45 +08:00
parent e6787a2036
commit 33c8720d69
5 changed files with 78 additions and 38 deletions
+15
View File
@@ -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.