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:
@@ -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."""
|
||||
|
||||
Reference in New Issue
Block a user