diff --git a/pipeline/io/export.py b/pipeline/io/export.py index 2061f42..8788c6b 100644 --- a/pipeline/io/export.py +++ b/pipeline/io/export.py @@ -83,6 +83,7 @@ def cache_jsonl( *, pack_size: int = -1, pad_value: int = 0, + batch_size: int = 256, ) -> List[str]: """Tokenize JSONL files and pack them into HDF5 storage. @@ -92,6 +93,7 @@ def cache_jsonl( processor: Initialized Processor instance. pack_size: Packing length, <=0 means no packing. pad_value: Padding value. + batch_size: Number of records passed to the processor at once. Returns: List of generated H5 file paths. @@ -105,25 +107,48 @@ def cache_jsonl( arrows: Dict[str, List] = {key: [] for key in output_keys} + def append_batch(batch): + items = [item for _, item in batch] + try: + results = processor.process_batch(items) + if len(results) != len(items): + raise RuntimeError( + "Batch processor returned a different number of results" + ) + except Exception: + results = [] + for line_num, item in batch: + 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: + arrows[key].append(result[key]) + + batch = [] + 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 ): try: - result = processor.process(json.loads(line)) - if result is not None: - for key in output_keys: - arrows[key].append(result[key]) + batch.append((line_num, json.loads(line))) + if len(batch) >= batch_size: + append_batch(batch) + batch = [] 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 + if batch: + append_batch(batch) if pack_size > 0: dtypes = ( diff --git a/pipeline/processors/base.py b/pipeline/processors/base.py index fb0de2f..5c52996 100644 --- a/pipeline/processors/base.py +++ b/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]: diff --git a/pipeline/processors/dpo.py b/pipeline/processors/dpo.py index f0a4938..f173873 100644 --- a/pipeline/processors/dpo.py +++ b/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"] diff --git a/pipeline/processors/pretrain.py b/pipeline/processors/pretrain.py index f3be6f6..ceaea5e 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"{item['text']}{self._eos_token}" for item in input_dicts] + encoded = self.tokenizer.encode(texts) + return [ + {"sequence": torch.tensor(tokens, dtype=torch.int32)} + for tokens in encoded + ] + @property def output_keys(self) -> List[str]: return ["sequence"] diff --git a/pipeline/processors/sft.py b/pipeline/processors/sft.py index c11464c..c929770 100644 --- a/pipeline/processors/sft.py +++ b/pipeline/processors/sft.py @@ -29,15 +29,20 @@ class SFTProcessor(BaseProcessor): - sequence: int32 tensor - Combined token IDs (prompt + response) - loss_mask: bool tensor - True for response tokens (compute loss) - 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: @@ -63,6 +68,35 @@ class SFTProcessor(BaseProcessor): "Input must contain 'messages' or 'query'/'response' pair" ) + def process_batch(self, input_dicts: List[Dict[str, Any]]) -> List[Dict[str, Tensor]]: + results: List[Optional[Dict[str, Tensor]]] = [None] * len(input_dicts) + message_indices = [i for i, item in enumerate(input_dicts) if "messages" in item] + legacy_indices = [ + i + for i, item in enumerate(input_dicts) + if "messages" not in item and "query" in item and "response" in item + ] + if len(message_indices) + len(legacy_indices) != len(input_dicts): + raise KeyError("Input must contain 'messages' or 'query'/'response' pair") + + if message_indices: + items = [input_dicts[i] for i in message_indices] + batch_results = self._process_messages_batch( + [item["messages"] for item in items] + ) + for index, result in zip(message_indices, batch_results): + results[index] = result + + if legacy_indices: + items = [input_dicts[i] for i in legacy_indices] + batch_results = self._process_legacy_batch(items) + for index, result in zip(legacy_indices, batch_results): + results[index] = result + + if any(result is None for result in results): + raise RuntimeError("Batch processing did not produce all results") + return results + def _process_messages(self, messages: List[Dict[str, str]]) -> Dict[str, Tensor]: if not messages: raise ValueError("Messages list is empty") @@ -73,21 +107,80 @@ class SFTProcessor(BaseProcessor): i for i, m in enumerate(messages) if m["role"] == "assistant" ) - prompt_tokens = self.tokenizer.apply_chat_template( + full_text = self.tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=False + ) + full_ids = self.tokenizer.encode(full_text, add_special_tokens=False) + + prompt_text = self.tokenizer.apply_chat_template( messages[:last_asst_idx], + tokenize=False, add_generation_prompt=True, - tokenize=True, ) + prompt_ids = self.tokenizer.encode(prompt_text, add_special_tokens=False) - resp_content = messages[last_asst_idx]["content"] - im_end = getattr(self.tokenizer, "im_end", "<|im_end|>") - resp_tokens = self.tokenizer.encode( - f"{resp_content}{im_end}\n", add_special_tokens=False - ) + resp_ids = full_ids[len(prompt_ids) :] + if not resp_ids: + raise ValueError("Empty assistant response") + + tokens, loss_mask = encode_with_mask(prompt_ids, list(resp_ids)) + + if self.max_seq_len and len(tokens) > self.max_seq_len: + tokens = tokens[: self.max_seq_len] + loss_mask = loss_mask[: self.max_seq_len] - tokens, loss_mask = encode_with_mask(prompt_tokens, resp_tokens) position_ids = torch.arange(len(tokens), dtype=torch.int32) - return {"sequence": tokens, "loss_mask": loss_mask, "position_ids": position_ids} + return { + "sequence": tokens, + "loss_mask": loss_mask, + "position_ids": position_ids, + } + + def _process_messages_batch( + self, conversations: List[List[Dict[str, str]]] + ) -> List[Dict[str, Tensor]]: + for messages in conversations: + if not messages: + raise ValueError("Messages list is empty") + if messages[-1]["role"] != "assistant": + raise ValueError("Last message must have role 'assistant'") + + assistant_indices = [ + max(i for i, message in enumerate(messages) if message["role"] == "assistant") + for messages in conversations + ] + full_texts = [ + self.tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=False + ) + for messages in conversations + ] + prompt_texts = [ + self.tokenizer.apply_chat_template( + messages[:assistant_idx], tokenize=False, add_generation_prompt=True + ) + for messages, assistant_idx in zip(conversations, assistant_indices) + ] + full_ids_batch = self.tokenizer.encode(full_texts, add_special_tokens=False) + prompt_ids_batch = self.tokenizer.encode(prompt_texts, add_special_tokens=False) + + results = [] + for full_ids, prompt_ids in zip(full_ids_batch, prompt_ids_batch): + resp_ids = full_ids[len(prompt_ids) :] + if not resp_ids: + raise ValueError("Empty assistant response") + tokens, loss_mask = encode_with_mask(prompt_ids, list(resp_ids)) + if self.max_seq_len and len(tokens) > self.max_seq_len: + tokens = tokens[: self.max_seq_len] + loss_mask = loss_mask[: self.max_seq_len] + results.append( + { + "sequence": tokens, + "loss_mask": loss_mask, + "position_ids": torch.arange(len(tokens), dtype=torch.int32), + } + ) + return results def _process_legacy(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]: strategy = self.strategy or ChatMLStrategy(self.tokenizer) @@ -102,6 +195,28 @@ class SFTProcessor(BaseProcessor): position_ids = torch.arange(len(tokens), dtype=torch.int32) return {"sequence": tokens, "loss_mask": loss_mask, "position_ids": position_ids} + def _process_legacy_batch( + self, input_dicts: List[Dict[str, Any]] + ) -> List[Dict[str, Tensor]]: + strategy = self.strategy or ChatMLStrategy(self.tokenizer) + query_batch = self.tokenizer.encode([item["query"] for item in input_dicts]) + response_batch = self.tokenizer.encode( + [item["response"] for item in input_dicts] + ) + results = [] + for query_tokens, response_tokens in zip(query_batch, response_batch): + prompt = strategy.assemble_prompt(query_tokens) + response = strategy.assemble_response(response_tokens) + tokens, loss_mask = encode_with_mask(prompt, response) + results.append( + { + "sequence": tokens, + "loss_mask": loss_mask, + "position_ids": torch.arange(len(tokens), dtype=torch.int32), + } + ) + return results + @property def output_keys(self) -> List[str]: return ["sequence", "loss_mask", "position_ids"] diff --git a/pipeline/tokenize/chat_template.py b/pipeline/tokenize/chat_template.py index 2922cc3..c59fa84 100644 --- a/pipeline/tokenize/chat_template.py +++ b/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 jinja_template.render(**variables) + return self._compiled.render(**variables) # Default ChatML template diff --git a/pipeline/tokenize/tokenizer.py b/pipeline/tokenize/tokenizer.py index 83f0f63..69ff159 100644 --- a/pipeline/tokenize/tokenizer.py +++ b/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 jinja_template.render(**variables) + return self._compiled.render(**variables) @@ -326,7 +330,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: diff --git a/scripts/cache_h5.py b/scripts/cache_h5.py index 8dc8899..b321686 100644 --- a/scripts/cache_h5.py +++ b/scripts/cache_h5.py @@ -53,6 +53,12 @@ def main(): parser.add_argument( "--pad-value", type=int, default=0, help="Padding value (default: 0)" ) + parser.add_argument( + "--batch-size", + type=int, + default=256, + help="Records tokenized per batch (default: 256)", + ) parser.add_argument( "--log-level", default="INFO", @@ -105,6 +111,7 @@ def main(): processor=processor, pack_size=args.pack_size, pad_value=args.pad_value, + batch_size=args.batch_size, ) print(f"\nDone! Output saved to {output_dir}") diff --git a/tests/test_cache.py b/tests/test_cache.py index c2c8585..f29f59b 100644 --- a/tests/test_cache.py +++ b/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) == 1 + + 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] diff --git a/tests/test_processors.py b/tests/test_processors.py index e10ba3d..d4e3ea0 100644 --- a/tests/test_processors.py +++ b/tests/test_processors.py @@ -15,7 +15,9 @@ from pipeline.processors import ( class DummyTokenizer: im_end = "<|im_end|>" - 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 apply_chat_template( @@ -50,6 +52,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): @@ -121,6 +133,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": []}) @@ -173,6 +202,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):