ViperEkura/DataPipeline
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: processors 支持批量 tokenize，优化性能并缓存 chat template

Browse Source
This commit is contained in:
ViperEkura
2026-08-05 12:21:25 +08:00
parent aa2ea4f3a6
commit 65dadac10f
10 changed files with 290 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pipeline/io/export.py
+35 -10
View File
@@ -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 = (
pipeline/processors/base.py
+6
View File
@@ -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
+27
View File
@@ -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/pretrain.py
+8
View File
@@ -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"]
pipeline/processors/sft.py
+124 -9
View File
@@ -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"]
pipeline/tokenize/chat_template.py
+6 -2
View File
@@ -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
pipeline/tokenize/tokenizer.py
+9 -3
View File
@@ -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:
scripts/cache_h5.py
+7
View File
@@ -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}")
tests/test_cache.py
+26
View File
@@ -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]
tests/test_processors.py
+42 -1
View File
@@ -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):