feat : preprocessing 支持 DPO/GRPO 多输出格式
- InputConfig 新增 sources 字段驱动多输出映射 - SectionedMaskBuilder 提取 _process_sections/_build_multi 模板方法 - Pipeline 泛化 accumulate 逻辑处理多 key 结果 - 测试拆分为 config/builder/pipeline 三文件，纯函数风格
This commit is contained in:
@@ -1,4 +1,9 @@
|
||||
"""Pipeline configuration for JSONL preprocessing."""
|
||||
"""Pipeline configuration for JSONL preprocessing.
|
||||
|
||||
Supports single-sequence (SFT/pretrain) and multi-output (DPO/GRPO)
|
||||
modes, both driven declaratively through ``input.sections`` or
|
||||
``input.sources``.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Dict, List, Optional
|
||||
@@ -8,7 +13,22 @@ from astrai.config.base import BaseConfig
|
||||
|
||||
@dataclass
|
||||
class InputConfig(BaseConfig):
|
||||
"""Declarative input mapping.
|
||||
|
||||
Single-output mode (backward-compatible)::
|
||||
|
||||
{"input": {"sections": [{"field": "messages", ...}]}}
|
||||
|
||||
Multi-output mode (DPO / GRPO)::
|
||||
|
||||
{"input": {"sources": {
|
||||
"chosen": {"sections": [{"field": "chosen", ...}]},
|
||||
"rejected": {"sections": [{"field": "rejected", ...}]},
|
||||
}}}
|
||||
"""
|
||||
|
||||
sections: Optional[List[Dict]] = None
|
||||
sources: Optional[Dict[str, Dict]] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
|
||||
+238
-59
@@ -1,7 +1,8 @@
|
||||
"""Mask building strategies for preprocessing pipeline.
|
||||
|
||||
The single :class:`SectionedMaskBuilder` handles all input formats
|
||||
via declarative ``input.sections`` config.
|
||||
(single-sequence / DPO / GRPO) via declarative config: ``input.sections``
|
||||
for single-output or ``input.sources`` for multi-output.
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
@@ -51,43 +52,142 @@ def _resolve_action(action: str, role: str, config) -> str:
|
||||
|
||||
@MaskBuilderFactory.register("sectioned")
|
||||
class SectionedMaskBuilder(BaseMaskBuilder):
|
||||
"""Config-driven builder: iterates over ``input.sections`` in order.
|
||||
"""Config-driven builder supporting single and multi-output modes.
|
||||
|
||||
Each section specifies a JSONL field + mask action.
|
||||
Single-output (backward-compatible)::
|
||||
|
||||
Section spec::
|
||||
|
||||
{
|
||||
"field": "messages", # JSONL key
|
||||
"action": "$role", # "train" | "mask" | "$role"
|
||||
"template": true, # apply chat_template per message (optional)
|
||||
"add_special_tokens": false # override encode flag (optional)
|
||||
}
|
||||
|
||||
Example configs::
|
||||
|
||||
# Chat
|
||||
{"input": {"sections": [
|
||||
{"field": "messages", "action": "$role", "template": true}
|
||||
]}}
|
||||
→ {"sequence": [...], "loss_mask": [...], "domain": "..."}
|
||||
|
||||
# Instruction
|
||||
{"input": {"sections": [
|
||||
{"field": "prompt", "action": "mask", "add_special_tokens": true},
|
||||
{"field": "response", "action": "train"}
|
||||
]}}
|
||||
Multi-output (DPO / GRPO)::
|
||||
|
||||
# Text
|
||||
{"input": {"sections": [
|
||||
{"field": "text", "action": "train"}
|
||||
]}}
|
||||
{"input": {"sources": {
|
||||
"chosen": {"sections": [
|
||||
{"field": "chosen", "action": "$role", "template": true}
|
||||
]},
|
||||
"rejected": {"sections": [
|
||||
{"field": "rejected", "action": "$role", "template": true}
|
||||
]}
|
||||
}}}
|
||||
→ {"chosen": [...], "chosen_mask": [...],
|
||||
"rejected": [...], "rejected_mask": [...], "domain": "..."}
|
||||
|
||||
Output spec fields::
|
||||
|
||||
sections – list of section specs (same format as single-output)
|
||||
list_field – True when the JSONL field holds a list of values to
|
||||
tokenise individually and concatenate (GRPO responses)
|
||||
mask_key – explicit output key for the loss mask
|
||||
(default: ``"{output_key}_mask"``)
|
||||
dtype – explicit tensor dtype for this output key
|
||||
(default: "int32")
|
||||
"""
|
||||
|
||||
def build(self, item: dict, config, tokenizer) -> Optional[dict]:
|
||||
sources_spec = getattr(config.input, "sources", None)
|
||||
if sources_spec:
|
||||
return self._build_multi(item, sources_spec, config, tokenizer)
|
||||
return self._build_single(item, config, tokenizer)
|
||||
|
||||
def _build_single(self, item: dict, config, tokenizer) -> Optional[dict]:
|
||||
sections = config.input.sections
|
||||
if not sections:
|
||||
return None
|
||||
|
||||
ids, mask = self._process_sections(
|
||||
item, sections, config, tokenizer, is_top_level=True
|
||||
)
|
||||
if ids is None:
|
||||
return None
|
||||
|
||||
result: dict = {
|
||||
"sequence": ids,
|
||||
"domain": _extract_domain(item, config.output.domain_key),
|
||||
}
|
||||
if not all(m == 1 for m in mask):
|
||||
result["loss_mask"] = mask
|
||||
return result
|
||||
|
||||
def _build_multi(
|
||||
self, item: dict, sources_spec: dict, config, tokenizer
|
||||
) -> Optional[dict]:
|
||||
result: dict = {}
|
||||
any_output = False
|
||||
|
||||
for output_key, spec in sources_spec.items():
|
||||
sections = spec.get("sections", [])
|
||||
if not sections:
|
||||
continue
|
||||
|
||||
if self._is_value_section(sections):
|
||||
ids = self._extract_raw_value(item, sections)
|
||||
if ids is None:
|
||||
continue
|
||||
result[output_key] = ids
|
||||
any_output = True
|
||||
continue
|
||||
|
||||
list_field = spec.get("list_field", False)
|
||||
mask_key = spec.get("mask_key", f"{output_key}_mask")
|
||||
|
||||
if list_field:
|
||||
ids, mask = self._process_list_field(item, sections, config, tokenizer)
|
||||
else:
|
||||
ids, mask = self._process_sections(
|
||||
item, sections, config, tokenizer, is_top_level=True
|
||||
)
|
||||
|
||||
if ids is None:
|
||||
continue
|
||||
|
||||
result[output_key] = ids
|
||||
if not all(m == 1 for m in mask):
|
||||
result[mask_key] = mask
|
||||
elif "mask_key" in spec:
|
||||
result[mask_key] = mask
|
||||
|
||||
any_output = True
|
||||
|
||||
if not any_output:
|
||||
return None
|
||||
|
||||
result["domain"] = _extract_domain(item, config.output.domain_key)
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _is_value_section(sections: list) -> bool:
|
||||
return len(sections) == 1 and sections[0].get("action") == "value"
|
||||
|
||||
@staticmethod
|
||||
def _extract_raw_value(item: dict, sections: list):
|
||||
"""Extract a raw value from a JSONL field without tokenisation.
|
||||
|
||||
Used for GRPO rewards where the field contains float values.
|
||||
"""
|
||||
sec = sections[0]
|
||||
field = sec["field"]
|
||||
raw = item.get(field)
|
||||
if raw is None:
|
||||
return None
|
||||
if isinstance(raw, list):
|
||||
return [float(v) for v in raw]
|
||||
return [float(raw)]
|
||||
|
||||
def _process_sections(
|
||||
self,
|
||||
item: dict,
|
||||
sections: list,
|
||||
config,
|
||||
tokenizer,
|
||||
*,
|
||||
is_top_level: bool = False,
|
||||
):
|
||||
"""Process a list of sections into ``(ids, loss_mask)``.
|
||||
|
||||
Returns ``(None, None)`` if the item should be skipped.
|
||||
"""
|
||||
all_ids: list[int] = []
|
||||
loss_mask: list[int] = []
|
||||
|
||||
@@ -96,7 +196,7 @@ class SectionedMaskBuilder(BaseMaskBuilder):
|
||||
s["action"] == "train" for s in sections
|
||||
)
|
||||
|
||||
if has_template and tokenizer.bos_token_id is not None:
|
||||
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)
|
||||
|
||||
@@ -110,33 +210,25 @@ class SectionedMaskBuilder(BaseMaskBuilder):
|
||||
)
|
||||
|
||||
if use_template:
|
||||
messages = item.get(field)
|
||||
if not isinstance(messages, list) or not messages:
|
||||
success = self._append_template_section(
|
||||
item, field, action, tokenizer, config, all_ids, loss_mask
|
||||
)
|
||||
if not success:
|
||||
continue
|
||||
for msg in messages:
|
||||
role = msg.get("role", "")
|
||||
act = _resolve_action(action, role, config)
|
||||
rendered = tokenizer.apply_chat_template(
|
||||
[msg], tokenize=False, add_generation_prompt=False
|
||||
)
|
||||
ids = tokenizer.encode(rendered, add_special_tokens=False)
|
||||
all_ids.extend(ids)
|
||||
val = 1 if act == "train" else 0
|
||||
loss_mask.extend([val] * len(ids))
|
||||
else:
|
||||
text = str(item.get(field, ""))
|
||||
if not text.strip():
|
||||
success = self._append_text_section(
|
||||
item,
|
||||
field,
|
||||
action,
|
||||
tokenizer,
|
||||
add_special,
|
||||
is_text_config,
|
||||
config,
|
||||
all_ids,
|
||||
loss_mask,
|
||||
)
|
||||
if not success:
|
||||
continue
|
||||
if is_text_config:
|
||||
pp = config.preprocessing
|
||||
if pp.min_chars > 0 and len(text) < pp.min_chars:
|
||||
continue
|
||||
if len(text) > pp.max_chars:
|
||||
continue
|
||||
ids = tokenizer.encode(text, add_special_tokens=add_special)
|
||||
all_ids.extend(ids)
|
||||
val = 1 if action == "train" else 0
|
||||
loss_mask.extend([val] * len(ids))
|
||||
|
||||
first_section = False
|
||||
|
||||
@@ -145,15 +237,102 @@ class SectionedMaskBuilder(BaseMaskBuilder):
|
||||
loss_mask = loss_mask[: len(all_ids)]
|
||||
|
||||
if not all_ids:
|
||||
return None
|
||||
return None, None
|
||||
|
||||
if has_template and len(all_ids) <= 1:
|
||||
return None
|
||||
if is_top_level and has_template and len(all_ids) <= 1:
|
||||
return None, None
|
||||
|
||||
result: dict = {
|
||||
"sequence": all_ids,
|
||||
"domain": _extract_domain(item, config.output.domain_key),
|
||||
}
|
||||
if not all(m == 1 for m in loss_mask):
|
||||
result["loss_mask"] = loss_mask
|
||||
return result
|
||||
return all_ids, loss_mask
|
||||
|
||||
def _append_template_section(
|
||||
self, item, field, action, tokenizer, config, all_ids, loss_mask
|
||||
):
|
||||
messages = item.get(field)
|
||||
if not isinstance(messages, list) or not messages:
|
||||
return False
|
||||
for msg in messages:
|
||||
role = msg.get("role", "")
|
||||
act = _resolve_action(action, role, config)
|
||||
rendered = tokenizer.apply_chat_template(
|
||||
[msg], tokenize=False, add_generation_prompt=False
|
||||
)
|
||||
ids = tokenizer.encode(rendered, add_special_tokens=False)
|
||||
all_ids.extend(ids)
|
||||
val = 1 if act == "train" else 0
|
||||
loss_mask.extend([val] * len(ids))
|
||||
return True
|
||||
|
||||
def _append_text_section(
|
||||
self,
|
||||
item,
|
||||
field,
|
||||
action,
|
||||
tokenizer,
|
||||
add_special,
|
||||
is_text_config,
|
||||
config,
|
||||
all_ids,
|
||||
loss_mask,
|
||||
):
|
||||
text = str(item.get(field, ""))
|
||||
if not text.strip():
|
||||
return False
|
||||
if is_text_config:
|
||||
pp = config.preprocessing
|
||||
if pp.min_chars > 0 and len(text) < pp.min_chars:
|
||||
return False
|
||||
if len(text) > pp.max_chars:
|
||||
return False
|
||||
ids = tokenizer.encode(text, add_special_tokens=add_special)
|
||||
all_ids.extend(ids)
|
||||
val = 1 if action == "train" else 0
|
||||
loss_mask.extend([val] * len(ids))
|
||||
return True
|
||||
|
||||
def _process_list_field(self, item: dict, sections: list, config, tokenizer):
|
||||
all_ids: list[int] = []
|
||||
loss_mask: list[int] = []
|
||||
|
||||
for sec in sections:
|
||||
field = sec["field"]
|
||||
action = sec["action"]
|
||||
use_template = sec.get("template", False)
|
||||
|
||||
values = item.get(field)
|
||||
if not isinstance(values, list):
|
||||
continue
|
||||
|
||||
for val in values:
|
||||
if use_template:
|
||||
if isinstance(val, list):
|
||||
wrapper = {field: val}
|
||||
self._append_template_section(
|
||||
wrapper,
|
||||
field,
|
||||
action,
|
||||
tokenizer,
|
||||
config,
|
||||
all_ids,
|
||||
loss_mask,
|
||||
)
|
||||
else:
|
||||
wrapper = {field: str(val)}
|
||||
self._append_text_section(
|
||||
wrapper,
|
||||
field,
|
||||
action,
|
||||
tokenizer,
|
||||
False,
|
||||
False,
|
||||
config,
|
||||
all_ids,
|
||||
loss_mask,
|
||||
)
|
||||
|
||||
max_len = config.preprocessing.max_seq_len
|
||||
all_ids = all_ids[:max_len]
|
||||
loss_mask = loss_mask[: len(all_ids)]
|
||||
|
||||
if not all_ids:
|
||||
return None, None
|
||||
return all_ids, loss_mask
|
||||
|
||||
@@ -81,17 +81,20 @@ class Pipeline:
|
||||
if result is None:
|
||||
continue
|
||||
|
||||
ids = result.pop("sequence")
|
||||
domain = result.pop("domain", "__default__")
|
||||
|
||||
is_multi = bool(getattr(self.config.input, "sources", None))
|
||||
if is_multi:
|
||||
ids = self._primary_ids(result)
|
||||
else:
|
||||
ids = result.pop("sequence")
|
||||
result["sequence"] = ids
|
||||
|
||||
if not ids:
|
||||
continue
|
||||
|
||||
domain = result.pop("domain", "__default__")
|
||||
result["sequence"] = ids
|
||||
|
||||
bucket = domains[domain]
|
||||
for key in list(bucket.keys()):
|
||||
if key not in result:
|
||||
bucket[key].append([1] * len(ids))
|
||||
self._align_bucket(bucket, result, ids, is_multi)
|
||||
for key, val in result.items():
|
||||
bucket[key].append(val)
|
||||
|
||||
@@ -108,6 +111,27 @@ class Pipeline:
|
||||
|
||||
print(f"Done. {count} documents tokenized.")
|
||||
|
||||
@staticmethod
|
||||
def _primary_ids(result: dict) -> list:
|
||||
"""Return the first list-valued entry in *result* as the primary id
|
||||
sequence for token counting."""
|
||||
for val in result.values():
|
||||
if isinstance(val, list) and val and isinstance(val[0], int):
|
||||
return val
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
def _align_bucket(bucket: dict, result: dict, ids: list, is_multi: bool):
|
||||
"""Pad previously-accumulated keys that are missing from *result*."""
|
||||
for key in list(bucket.keys()):
|
||||
if key in result:
|
||||
continue
|
||||
if is_multi:
|
||||
pad = bucket[key][-1] if bucket[key] else [1] * len(ids)
|
||||
bucket[key].append(pad)
|
||||
else:
|
||||
bucket[key].append([1] * len(ids))
|
||||
|
||||
def _iter_items(self):
|
||||
for path in self.paths:
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
@@ -135,7 +159,8 @@ class Pipeline:
|
||||
else:
|
||||
save_h5(chunk_dir, f"data_{idx:04d}", tensors)
|
||||
shard_idx[domain] = idx + 1
|
||||
first_key = "sequence" if "sequence" in tensors else next(iter(tensors))
|
||||
tqdm.tqdm.write(
|
||||
f" saved {domain}/shard_{idx:04d} "
|
||||
f"({tensors['sequence'][0].numel():,} tokens)"
|
||||
f"({tensors[first_key][0].numel():,} tokens)"
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user