From 33c8720d693680cf99aa14d01ce1dbdec676f9ba Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sat, 25 Jul 2026 12:19:45 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E5=B9=B6=E8=A1=8C=20batch=20tokenizati?= =?UTF-8?q?on=20+=20cache=5Fjsonl=20=E6=89=B9=E5=A4=84=E7=90=86=E6=94=AF?= =?UTF-8?q?=E6=8C=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 | 58 ++++++++++++++++++++------------- pipeline/processors/base.py | 15 +++++++++ pipeline/processors/pretrain.py | 8 +++++ pipeline/tokenize/tokenizer.py | 28 +++++++--------- scripts/cache_h5.py | 7 ++++ 5 files changed, 78 insertions(+), 38 deletions(-) diff --git a/pipeline/io/export.py b/pipeline/io/export.py index ca978e1..d077099 100644 --- a/pipeline/io/export.py +++ b/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: diff --git a/pipeline/processors/base.py b/pipeline/processors/base.py index fb0de2f..4e2e859 100644 --- a/pipeline/processors/base.py +++ b/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. diff --git a/pipeline/processors/pretrain.py b/pipeline/processors/pretrain.py index f3be6f6..4d3aae7 100644 --- a/pipeline/processors/pretrain.py +++ b/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"] diff --git a/pipeline/tokenize/tokenizer.py b/pipeline/tokenize/tokenizer.py index 060adbe..50d91df 100644 --- a/pipeline/tokenize/tokenizer.py +++ b/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.""" diff --git a/scripts/cache_h5.py b/scripts/cache_h5.py index fcae8c7..a75db88 100644 --- a/scripts/cache_h5.py +++ b/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}")