ViperEkura
33c8720d69
- 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
57 lines
1.6 KiB
Python
57 lines
1.6 KiB
Python
"""Pre-training data processor."""
|
||
|
||
from typing import Any, Dict, List
|
||
|
||
import torch
|
||
from torch import Tensor
|
||
|
||
from pipeline.tokenize import AutoTokenizer
|
||
from pipeline.processors.base import BaseProcessor, ProcessorSchema
|
||
from pipeline.processors.factory import ProcessorFactory
|
||
|
||
|
||
@ProcessorFactory.register("pt")
|
||
class PreTrainProcessor(BaseProcessor):
|
||
"""Pre-training data processor.
|
||
|
||
Processes raw text into tokenized sequences with EOS tokens.
|
||
|
||
Input schema:
|
||
- text: str - Raw text string to tokenize
|
||
|
||
Output schema:
|
||
- sequence: int32 tensor - Token IDs with EOS appended
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
tokenizer: AutoTokenizer,
|
||
eos_token: str = "<｜end▁of▁sentence｜>",
|
||
):
|
||
self.tokenizer = tokenizer
|
||
self._eos_token = eos_token
|
||
|
||
@property
|
||
def schema(self) -> ProcessorSchema:
|
||
return ProcessorSchema(
|
||
input_fields={"text": str},
|
||
output_fields={"sequence": torch.int32},
|
||
)
|
||
|
||
def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
|
||
segment = input_dict["text"]
|
||
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"]
|