From 99b5d2b2da61de6eeb9ba80482a2d8a71dc2d972 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Thu, 23 Jul 2026 18:42:19 +0800 Subject: [PATCH] perf: batch tokenizer preprocessing --- astrai/config/preprocess_config.py | 3 + astrai/preprocessing/builder.py | 200 ++++++++++++++++++++++++++ astrai/preprocessing/pipeline.py | 80 +++++++---- scripts/tools/preprocess.py | 10 ++ tests/data/test_preprocess_builder.py | 52 +++++++ 5 files changed, 320 insertions(+), 25 deletions(-) diff --git a/astrai/config/preprocess_config.py b/astrai/config/preprocess_config.py index f10070c..91219d4 100644 --- a/astrai/config/preprocess_config.py +++ b/astrai/config/preprocess_config.py @@ -45,6 +45,8 @@ class ProcessingConfig(BaseConfig): Maximum number of characters to keep (default: 2_000_000). max_items : Optional[int] Maximum number of items to process (default: None, unlimited). + batch_size : int + Number of records tokenized together (default: 256). packing_strategy : str How to pack sequences into a contiguous stream. @@ -65,6 +67,7 @@ class ProcessingConfig(BaseConfig): min_chars: int = 50 max_chars: int = 2_000_000 max_items: Optional[int] = None + batch_size: int = 256 packing_strategy: str = "simple" max_packed_len: int = 8192 truncation_mode: str = "keep_start" diff --git a/astrai/preprocessing/builder.py b/astrai/preprocessing/builder.py index 17da350..5d5778a 100644 --- a/astrai/preprocessing/builder.py +++ b/astrai/preprocessing/builder.py @@ -94,6 +94,97 @@ class SectionRenderer: return all_ids, loss_mask + def process_sections_batch( + self, + items: list[dict], + sections: list, + config, + tokenizer, + *, + is_top_level=False, + filter_text=True, + ): + """Render and tokenize a group of records with batched Rust tokenization.""" + has_template = any(s.get("template") for s in sections) + is_text_config = not has_template and all( + s["action"] == "train" for s in sections + ) + plans: list[list[tuple[str, str, bool]]] = [] + + for item in items: + plan: list[tuple[str, str, bool]] = [] + first_section = True + for sec in sections: + field = sec["field"] + action = sec["action"] + use_template = sec.get("template", False) + add_special = sec.get( + "add_special_tokens", not use_template and first_section + ) + + if use_template: + messages = item.get(field) + if not isinstance(messages, list) or not messages: + continue + for msg in messages: + role = msg.get("role", "") + rendered = tokenizer.apply_chat_template( + [msg], tokenize=False, add_generation_prompt=False + ) + plan.append( + (rendered, _resolve_action(action, role, config), False) + ) + else: + text = str(item.get(field, "")) + if not text.strip(): + continue + if is_text_config and filter_text: + pp = config.preprocessing + if pp.min_chars > 0 and len(text) < pp.min_chars: + continue + if len(text) > pp.max_chars: + continue + plan.append((text, action, add_special)) + + first_section = False + plans.append(plan) + + encoded: dict[tuple[int, int], list[int]] = {} + for add_special in (False, True): + refs = [ + (item_idx, unit_idx, text) + for item_idx, plan in enumerate(plans) + for unit_idx, (text, _, add) in enumerate(plan) + if add == add_special + ] + if not refs: + continue + ids_batch = tokenizer.encode( + [text for _, _, text in refs], add_special_tokens=add_special + ) + for (item_idx, unit_idx, _), ids in zip(refs, ids_batch): + encoded[(item_idx, unit_idx)] = ids + + outputs = [] + max_len = config.preprocessing.max_seq_len + for item_idx, plan in enumerate(plans): + all_ids = [] + loss_mask = [] + if is_top_level and has_template and tokenizer.bos_token_id is not None: + all_ids.append(tokenizer.bos_token_id) + loss_mask.append(0) + for unit_idx, (_, action, _) in enumerate(plan): + ids = encoded[(item_idx, unit_idx)] + all_ids.extend(ids) + loss_mask.extend([1 if action == "train" else 0] * len(ids)) + all_ids = all_ids[:max_len] + loss_mask = loss_mask[: len(all_ids)] + if not all_ids or (is_top_level and has_template and len(all_ids) <= 1): + outputs.append((None, None)) + else: + outputs.append((all_ids, loss_mask)) + return outputs + def process_list_field(self, item: dict, sections: list, config, tokenizer): """Tokenize a list-valued field, preserving per-element boundaries. @@ -147,6 +238,42 @@ class SectionRenderer: return None, None return per_item_ids, per_item_masks + def process_list_field_batch(self, items, sections, config, tokenizer): + per_item_ids = [[] for _ in items] + per_item_masks = [[] for _ in items] + + for sec in sections: + wrappers = [] + owners = [] + field = sec["field"] + for item_idx, item in enumerate(items): + values = item.get(field) + if not isinstance(values, list): + continue + for val in values: + if sec.get("template", False) and not isinstance(val, list): + continue + wrappers.append({field: val if isinstance(val, list) else str(val)}) + owners.append(item_idx) + + rendered = self.process_sections_batch( + wrappers, + [sec], + config, + tokenizer, + is_top_level=False, + filter_text=False, + ) + for owner, (ids, mask) in zip(owners, rendered): + if ids: + per_item_ids[owner].append(ids) + per_item_masks[owner].append(mask) + + return [ + (ids, masks) if ids else (None, None) + for ids, masks in zip(per_item_ids, per_item_masks) + ] + @staticmethod def is_value_section(sections: list) -> bool: return len(sections) == 1 and sections[0].get("action") == "value" @@ -214,6 +341,9 @@ class BaseMaskBuilder(ABC): @abstractmethod def build(self, item: dict, config, tokenizer) -> Optional[dict]: ... + def build_batch(self, items: list[dict], config, tokenizer) -> list[Optional[dict]]: + return [self.build(item, config, tokenizer) for item in items] + class MaskBuilderFactory(BaseFactory["BaseMaskBuilder"]): pass @@ -248,6 +378,27 @@ class SingleOutputMaskBuilder(BaseMaskBuilder): result["loss_mask"] = mask return result + def build_batch(self, items, config, tokenizer): + sections = config.input.sections + if not sections: + return [None] * len(items) + rendered = self.renderer.process_sections_batch( + items, sections, config, tokenizer, is_top_level=True + ) + results = [] + for item, (ids, mask) in zip(items, rendered): + if ids is None: + results.append(None) + continue + result = { + "sequence": ids, + "domain": _extract_domain(item, config.output.domain_key), + } + if not all(m == 1 for m in mask): + result["loss_mask"] = mask + results.append(result) + return results + @MaskBuilderFactory.register("multi") class MultiOutputMaskBuilder(BaseMaskBuilder): @@ -317,6 +468,49 @@ class MultiOutputMaskBuilder(BaseMaskBuilder): result["domain"] = _extract_domain(item, config.output.domain_key) return result + def build_batch(self, items, config, tokenizer): + sources_spec = getattr(config.input, "sources", None) + if not sources_spec: + return [None] * len(items) + + results = [{} for _ in items] + for output_key, spec in sources_spec.items(): + sections = spec.get("sections", []) + if not sections: + continue + if self.renderer.is_value_section(sections): + for item, result in zip(items, results): + value = self.renderer.extract_raw_value(item, sections) + if value is not None: + result[output_key] = value + continue + + mask_key = spec.get("mask_key", f"{output_key}_mask") + if spec.get("list_field", False): + rendered = self.renderer.process_list_field_batch( + items, sections, config, tokenizer + ) + else: + rendered = self.renderer.process_sections_batch( + items, sections, config, tokenizer, is_top_level=True + ) + + for result, (ids, mask) in zip(results, rendered): + if ids is None: + continue + result[output_key] = ids + if spec.get("list_field", False) or not all(m == 1 for m in mask): + result[mask_key] = mask + elif "mask_key" in spec: + result[mask_key] = mask + + return [ + ({**result, "domain": _extract_domain(item, config.output.domain_key)}) + if result + else None + for item, result in zip(items, results) + ] + @MaskBuilderFactory.register("sectioned") class SectionedMaskBuilder(BaseMaskBuilder): @@ -335,3 +529,9 @@ class SectionedMaskBuilder(BaseMaskBuilder): if sources_spec: return self._multi.build(item, config, tokenizer) return self._single.build(item, config, tokenizer) + + def build_batch(self, items, config, tokenizer): + sources_spec = getattr(config.input, "sources", None) + if sources_spec: + return self._multi.build_batch(items, config, tokenizer) + return self._single.build_batch(items, config, tokenizer) diff --git a/astrai/preprocessing/pipeline.py b/astrai/preprocessing/pipeline.py index 3d15a3f..009bd41 100644 --- a/astrai/preprocessing/pipeline.py +++ b/astrai/preprocessing/pipeline.py @@ -80,6 +80,9 @@ class Pipeline: def transform(self, item: dict) -> Optional[dict]: return self.mask_builder.build(item, self.config, self.tokenizer) + def transform_batch(self, items: list[dict]) -> list[Optional[dict]]: + return self.mask_builder.build_batch(items, self.config, self.tokenizer) + def run(self): domains: dict = defaultdict(lambda: defaultdict(list)) total_tokens = 0 @@ -88,39 +91,55 @@ class Pipeline: pp = self.config.preprocessing - for item in tqdm.tqdm( - self._iter_items(), desc="Tokenizing", unit="docs", mininterval=0.5 - ): - if pp.max_items and count >= pp.max_items: - break - + progress = tqdm.tqdm(desc="Tokenizing", unit="docs", mininterval=0.5) + stop = False + for items in self._iter_batches(pp.batch_size): + progress.update(len(items)) try: - result = self.transform(item) + results = self.transform_batch(items) except Exception: logger.warning( - "Failed to process item #%d, skipping", count + 1, exc_info=True + "Failed to process batch, retrying records individually", + exc_info=True, ) - continue - if result is None: - continue + results = [] + for item in items: + try: + results.append(self.transform(item)) + except Exception: + logger.warning( + "Failed to process item, skipping", exc_info=True + ) + results.append(None) - domain = result.pop("domain", "__default__") - ids = primary_ids(result) - if not ids: - continue + for result in results: + if pp.max_items and count >= pp.max_items: + stop = True + break + if result is None: + continue - bucket = domains[domain] - self._align_bucket(bucket, result, ids) - for key, val in result.items(): - bucket[key].append(val) + domain = result.pop("domain", "__default__") + ids = primary_ids(result) + if not ids: + continue - count += 1 - total_tokens += len(ids) + bucket = domains[domain] + self._align_bucket(bucket, result, ids) + for key, val in result.items(): + bucket[key].append(val) - if total_tokens >= self.config.output.max_tokens_per_shard: - self._flush(domains, shard_idx) - domains.clear() - total_tokens = 0 + count += 1 + total_tokens += len(ids) + + if total_tokens >= self.config.output.max_tokens_per_shard: + self._flush(domains, shard_idx) + domains.clear() + total_tokens = 0 + if stop: + break + + progress.close() if total_tokens > 0: self._flush(domains, shard_idx) @@ -149,6 +168,17 @@ class Pipeline: continue yield json.loads(line) + def _iter_batches(self, batch_size: int): + batch_size = max(1, batch_size) + batch = [] + for item in self._iter_items(): + batch.append(item) + if len(batch) >= batch_size: + yield batch + batch = [] + if batch: + yield batch + def _flush(self, domains, shard_idx): for domain, keys in domains.items(): idx = shard_idx[domain] diff --git a/scripts/tools/preprocess.py b/scripts/tools/preprocess.py index 4e19f42..0466d66 100644 --- a/scripts/tools/preprocess.py +++ b/scripts/tools/preprocess.py @@ -22,9 +22,19 @@ def main(): default="params", help="Path to tokenizer directory (default: params)", ) + parser.add_argument( + "--batch_size", + type=int, + default=None, + help="Number of records tokenized together (default: config value)", + ) args = parser.parse_args() config = PipelineConfig.from_file(args.config) + if args.batch_size is not None: + if args.batch_size < 1: + parser.error("--batch_size must be at least 1") + config.preprocessing.batch_size = args.batch_size Pipeline( config=config, diff --git a/tests/data/test_preprocess_builder.py b/tests/data/test_preprocess_builder.py index abc8aa7..4dfe327 100644 --- a/tests/data/test_preprocess_builder.py +++ b/tests/data/test_preprocess_builder.py @@ -68,6 +68,28 @@ def test_chat_mask_only_assistant(chat_tokenizer, builder): assert len(masked) > 0 +def test_chat_batch_matches_single(chat_tokenizer, builder): + config = make_chat_config() + items = [ + { + "messages": [ + {"role": "user", "content": "What is 2+2?"}, + {"role": "assistant", "content": "4"}, + ] + }, + { + "messages": [ + {"role": "system", "content": "Be concise."}, + {"role": "user", "content": "Say hello."}, + {"role": "assistant", "content": "Hello."}, + ] + }, + ] + batch = builder.build_batch(items, config, chat_tokenizer) + single = [builder.build(item, config, chat_tokenizer) for item in items] + assert batch == single + + @pytest.mark.parametrize( "mask_rules,mask_default,expect_nonzero", [ @@ -152,6 +174,17 @@ def test_instruction_basic(test_tokenizer, builder): assert len(result["sequence"]) == len(result["loss_mask"]) +def test_instruction_batch_matches_single(test_tokenizer, builder): + config = make_instruction_config() + items = [ + {"prompt": "Translate to French: Hello", "response": "Bonjour"}, + {"prompt": "Translate to German: Hello", "response": "Hallo"}, + ] + assert builder.build_batch(items, config, test_tokenizer) == [ + builder.build(item, config, test_tokenizer) for item in items + ] + + def test_instruction_prompt_masked(test_tokenizer, builder): config = make_instruction_config() item = {"prompt": "hello", "response": "world"} @@ -363,6 +396,25 @@ def test_grpo_basic(chat_tokenizer, builder): assert result["rewards"] == [1.0, 0.5, 0.8, 0.2] +def test_grpo_batch_matches_single(chat_tokenizer, builder): + config = make_grpo_config() + items = [ + { + "prompt": [{"role": "user", "content": "What is 2+2?"}], + "responses": ["4", "5"], + "rewards": [1.0, 0.0], + }, + { + "prompt": [{"role": "user", "content": "Say hello."}], + "responses": ["Hello", "Hi"], + "rewards": [1.0, 0.5], + }, + ] + assert builder.build_batch(items, config, chat_tokenizer) == [ + builder.build(item, config, chat_tokenizer) for item in items + ] + + def test_grpo_response_tokens_all_trained(chat_tokenizer, builder): config = make_grpo_config() item = {