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
+8
View File
@@ -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"]