merge remote main

pipeline/io/export.py

@@ -4,7 +4,7 @@ import json
 import logging
 import os
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import torch
 from datasets import Dataset
@@ -160,23 +160,46 @@ def cache_jsonl(
         arrows_batch: Dict[str, List] = {key: [] for key in output_keys}
         batch_tokens: int = 0
 
-        buf: List[str] = []
+        buf: List[Tuple[int, str]] = []
-        buf_num: int = 0
 
         def flush_buf():
             nonlocal batch_tokens
             if not buf:
                 return
             samples = []
-            for line in buf:
+            for line_num, line in buf:
                 try:
-                    samples.append(json.loads(line))
+                    samples.append((line_num, json.loads(line)))
                 except json.JSONDecodeError as e:
-                    logger.warning(f"JSON decode error, skipping: {e}")
+                    logger.warning(
+                        f"JSON decode error in {file_path} line {line_num}: "
+                        f"{e}. Skipping line."
+                    )
             buf.clear()
             if not samples:
                 return
-            results = processor.process_batch(samples) if hasattr(processor, "process_batch") else [processor.process(s) for s in samples]
+            items = [item for _, item in samples]
+            try:
+                results = (
+                    processor.process_batch(items)
+                    if hasattr(processor, "process_batch")
+                    else [processor.process(s) for s in items]
+                )
+                if len(results) != len(items):
+                    raise RuntimeError(
+                        "Batch processor returned a different number of results"
+                    )
+            except Exception:
+                results = []
+                for line_num, item in samples:
+                    try:
+                        results.append(processor.process(item))
+                    except Exception as e:
+                        logger.warning(
+                            f"Unexpected error processing line {line_num} "
+                            f"in {file_path}: {e}. Skipping line."
+                        )
+                        results.append(None)
             for result in results:
                 if result is not None:
                     for key in output_keys:
@@ -184,15 +207,23 @@ def cache_jsonl(
                     if target_tokens > 0:
                         batch_tokens += int(result[output_keys[0]].shape[0])
 
+        batch_size = max(1, batch_size)
         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
             ):
-                buf.append(line)
+                buf.append((line_num, 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)
+                        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] = []

pipeline/processors/base.py

@@ -76,6 +76,12 @@ class BaseProcessor(ABC):
         """
         pass
 
+    def process_batch(
+        self, input_dicts: List[Dict[str, Any]]
+    ) -> List[Dict[str, Tensor]]:
+        """Process a batch, falling back to the single-record implementation."""
+        return [self.process(input_dict) for input_dict in input_dicts]
+
     @property
     @abstractmethod
     def output_keys(self) -> List[str]:

pipeline/processors/dpo.py

@@ -74,6 +74,33 @@ class DPOProcessor(BaseProcessor):
             "rejected_mask": rejected_m,
         }
 
+    def process_batch(self, input_dicts: List[Dict[str, Any]]) -> List[Dict[str, Tensor]]:
+        query_batch = self.tokenizer.encode([item["query"] for item in input_dicts])
+        chosen_batch = self.tokenizer.encode([item["chosen"] for item in input_dicts])
+        rejected_batch = self.tokenizer.encode(
+            [item["rejected"] for item in input_dicts]
+        )
+        results = []
+        for query_tokens, chosen_tokens, rejected_tokens in zip(
+            query_batch, chosen_batch, rejected_batch
+        ):
+            prompt = self.strategy.assemble_prompt(query_tokens)
+            chosen_t, chosen_m = encode_with_mask(
+                prompt, self.strategy.assemble_response(chosen_tokens)
+            )
+            rejected_t, rejected_m = encode_with_mask(
+                prompt, self.strategy.assemble_response(rejected_tokens)
+            )
+            results.append(
+                {
+                    "chosen": chosen_t,
+                    "chosen_mask": chosen_m,
+                    "rejected": rejected_t,
+                    "rejected_mask": rejected_m,
+                }
+            )
+        return results
+
     @property
     def output_keys(self) -> List[str]:
         return ["chosen", "chosen_mask", "rejected", "rejected_mask"]

pipeline/processors/sft.py

@@ -26,18 +26,23 @@ class SFTProcessor(BaseProcessor):
          Internally converted to messages.
 
     Output schema:
-        - sequence: int32 tensor - Combined token IDs
+        - sequence: int32 tensor - Combined token IDs (prompt + response)
-        - loss_mask: bool tensor - True for assistant response tokens
+        - loss_mask: bool tensor - True for response tokens (compute loss)
-        - position_ids: int32 tensor - Per-sample position IDs, start from 0
+        - position_ids: int32 tensor - Per-sample position IDs starting from 0
+
+    Only the final assistant message is trained (mask_history behavior).
+    All earlier turns are context/prompt and masked from loss.
     """
 
     def __init__(
         self,
         tokenizer: AutoTokenizer,
         strategy: Optional[PromptStrategy] = None,
+        max_seq_len: Optional[int] = None,
     ):
         self.tokenizer = tokenizer
         self.strategy = strategy
+        self.max_seq_len = max_seq_len
 
     @property
     def schema(self) -> ProcessorSchema:
@@ -89,6 +94,9 @@ class SFTProcessor(BaseProcessor):
         sequence = torch.tensor(prompt + resp, dtype=torch.int32)
         loss_mask = torch.zeros(len(sequence), dtype=torch.bool)
         loss_mask[len(prompt) :] = True
+        if self.max_seq_len and len(sequence) > self.max_seq_len:
+            sequence = sequence[: self.max_seq_len]
+            loss_mask = loss_mask[: self.max_seq_len]
         position_ids = torch.arange(len(sequence), dtype=torch.int32)
         return {
             "sequence": sequence,
@@ -135,6 +143,9 @@ class SFTProcessor(BaseProcessor):
             sequence = torch.tensor(prompt_tokens + resp_tokens, dtype=torch.int32)
             loss_mask = torch.zeros(len(sequence), dtype=torch.bool)
             loss_mask[len(prompt_tokens):] = True
+            if self.max_seq_len and len(sequence) > self.max_seq_len:
+                sequence = sequence[: self.max_seq_len]
+                loss_mask = loss_mask[: self.max_seq_len]
             position_ids = torch.arange(len(sequence), dtype=torch.int32)
             results[idx] = {
                 "sequence": sequence,

pipeline/tokenize/chat_template.py

@@ -3,6 +3,7 @@ Chat template module with Jinja2 rendering support.
 """
 
 from dataclasses import dataclass, field
+from functools import cached_property
 from typing import Any, Dict, List, Optional
 
 from jinja2 import Template
@@ -32,6 +33,10 @@ class ChatTemplate:
     default_variables: Dict[str, Any] = field(default_factory=dict)
     special_tokens: Dict[str, str] = field(default_factory=dict)
 
+    @cached_property
+    def _compiled(self) -> Template:
+        return Template(self.template_str)
+
     @classmethod
     def from_string(
         cls,
@@ -79,8 +84,7 @@ class ChatTemplate:
         if system_prompt is not None:
             variables["system_prompt"] = system_prompt
 
-        jinja_template = Template(self.template_str)
+        return self._compiled.render(**variables)
-        return jinja_template.render(**variables)
 
 
 # Default ChatML template

pipeline/tokenize/tokenizer.py

@@ -3,6 +3,7 @@ Tokenizer module with BPE implementation and auto-loading support.
 """
 
 from dataclasses import dataclass
+from functools import cached_property
 import json
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Union
@@ -102,6 +103,10 @@ class ChatTemplate:
         if self.special_tokens is None:
             self.special_tokens = {}
 
+    @cached_property
+    def _compiled(self) -> Template:
+        return Template(self.template_str)
+
     @classmethod
     def from_string(
         cls,
@@ -142,8 +147,7 @@ class ChatTemplate:
         if system_prompt is not None:
             variables["system_prompt"] = system_prompt
 
-        jinja_template = Template(self.template_str)
+        return self._compiled.render(**variables)
-        return jinja_template.render(**variables)
 
 
 
@@ -328,7 +332,9 @@ class AutoTokenizer:
             KeyError: If template name is not registered.
         """
         if isinstance(template, str):
-            self._chat_template = ChatTemplate.from_string(template)
+            self._chat_template = ChatTemplate.from_string(
+                template, special_tokens=self._special_token_map
+            )
         elif isinstance(template, ChatTemplate):
             self._chat_template = template
         else:

scripts/cache_h5.py

@@ -66,6 +66,12 @@ def main():
         default=1_000,
         help="Merge every N packed chunks into one tensor, <=0 to disable (default: 1000)",
     )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=256,
+        help="Records tokenized per batch (default: 256)",
+    )
     parser.add_argument(
         "--log-level",
         default="INFO",

tests/test_cache.py

@@ -29,6 +29,16 @@ class DummyProcessor(BaseProcessor):
         }
 
 
+class BatchTrackingProcessor(DummyProcessor):
+    def __init__(self):
+        super().__init__()
+        self.batch_sizes = []
+
+    def process_batch(self, items):
+        self.batch_sizes.append(len(items))
+        return super().process_batch(items)
+
+
 class TestCacheJsonl:
     def test_basic_cache_functionality(self):
         with tempfile.TemporaryDirectory() as tmpdir:
@@ -121,3 +131,19 @@ class TestCacheJsonl:
                 pad_value=0,
             )
             assert len(output_files) == 0
+
+    def test_uses_configured_batch_size(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            jsonl_path = os.path.join(tmpdir, "test.jsonl")
+            with open(jsonl_path, "w", encoding="utf-8") as f:
+                for i in range(5):
+                    f.write(json.dumps({"text": str(i)}) + "\n")
+
+            processor = BatchTrackingProcessor()
+            cache_jsonl(
+                files=[jsonl_path],
+                output_dir=tmpdir,
+                processor=processor,
+                batch_size=2,
+            )
+            assert processor.batch_sizes == [2, 2, 1]

tests/test_processors.py

@@ -17,7 +17,9 @@ class DummyTokenizer:
         self._special_token_map = {}
         self._chat_template = None
 
-    def encode(self, text: str, add_special_tokens: bool = False):
+    def encode(self, text, add_special_tokens: bool = False):
+        if isinstance(text, list):
+            return [[ord(c) for c in item] for item in text]
         return [ord(c) for c in text]
 
     def decode(self, tokens, skip_special_tokens=True):
@@ -59,6 +61,16 @@ class TestPreTrainProcessor:
         result = PreTrainProcessor(DummyTokenizer()).process({"text": "a"})
         assert len(result["sequence"]) > 0
 
+    def test_process_batch_matches_single(self):
+        processor = PreTrainProcessor(DummyTokenizer())
+        items = [{"text": "hello"}, {"text": "world"}]
+        batch = processor.process_batch(items)
+        single = [processor.process(item) for item in items]
+        assert all(
+            torch.equal(batch_item["sequence"], single_item["sequence"])
+            for batch_item, single_item in zip(batch, single)
+        )
+
 
 class TestSFTProcessor:
     def test_output_keys(self):
@@ -130,6 +142,23 @@ class TestSFTProcessor:
         })
         assert "sequence" in result
 
+    def test_process_batch_matches_single(self):
+        processor = SFTProcessor(DummyTokenizer())
+        items = [
+            {
+                "messages": [
+                    {"role": "user", "content": "q1"},
+                    {"role": "assistant", "content": "a1"},
+                ]
+            },
+            {"query": "q2", "response": "a2"},
+        ]
+        batch = processor.process_batch(items)
+        single = [processor.process(item) for item in items]
+        for batch_item, single_item in zip(batch, single):
+            for key in processor.output_keys:
+                assert torch.equal(batch_item[key], single_item[key])
+
     def test_messages_empty_raises(self):
         with pytest.raises(ValueError, match="Messages list is empty"):
             SFTProcessor(DummyTokenizer()).process({"messages": []})
@@ -182,6 +211,18 @@ class TestDPOProcessor:
         assert result["chosen_mask"].dtype == torch.bool
         assert result["rejected_mask"].dtype == torch.bool
 
+    def test_process_batch_matches_single(self):
+        processor = DPOProcessor(DummyTokenizer())
+        items = [
+            {"query": "q1", "chosen": "yes", "rejected": "no"},
+            {"query": "q2", "chosen": "good", "rejected": "bad"},
+        ]
+        batch = processor.process_batch(items)
+        single = [processor.process(item) for item in items]
+        for batch_item, single_item in zip(batch, single):
+            for key in processor.output_keys:
+                assert torch.equal(batch_item[key], single_item[key])
+
 
 class TestProcessorFactory:
     def test_create_pre_train_processor(self):