perf: batch tokenizer preprocessing

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"

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)

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,19 +91,31 @@ 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
+                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)
+
+            for result in results:
+                if pp.max_items and count >= pp.max_items:
+                    stop = True
+                    break
                 if result is None:
                     continue
 
@@ -121,6 +136,10 @@ class Pipeline:
                     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]

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,

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 = {