refactor: deduplicate preprocessing kernel and BFD packing
- Extract shared core (mask building, primary-id extraction, tensorisation, position-id generation) to astrai/preprocessing/core.py; Pipeline and TokenizeTransform both consume it, eliminating ~60% duplicated logic - Promote BFD _plan to module-level plan_bfd(lengths, max_len) returning pure index bins; BFDPacking.apply and evaluate_ifd._pack_bins both call it, removing the second BFD implementation - Split Pipeline._flush (49 lines) into _inject_doc_reset_position_ids + _inject_continuous_position_ids + _to_tensors; split Pipeline.run by delegating record iteration to core.iter_raw_records - Remove dead no-op pop/塞回 in Pipeline.run (L110-111)
This commit is contained in:
@@ -0,0 +1,124 @@
|
||||
"""Shared preprocessing kernel used by both :class:`Pipeline` and
|
||||
:class:`TokenizeTransform`.
|
||||
|
||||
The two entry points previously duplicated ~60 % of their logic:
|
||||
record iteration, mask-builder invocation, primary-id extraction,
|
||||
per-key accumulation, dtype inference and position-id generation.
|
||||
This module factors out the common core as pure functions so that
|
||||
the online (``TokenizeTransform``) and offline (``Pipeline``) paths
|
||||
stay in lockstep.
|
||||
"""
|
||||
|
||||
from itertools import chain
|
||||
from typing import Dict, Iterator, List, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from astrai.config.preprocess_config import PipelineConfig
|
||||
from astrai.preprocessing.builder import MaskBuilderFactory
|
||||
from astrai.preprocessing.position_id import PositionIdStrategyFactory
|
||||
from astrai.tokenize import AutoTokenizer
|
||||
|
||||
|
||||
def build_preprocessing_components(config: PipelineConfig, tokenizer_path: str):
|
||||
"""Load tokenizer, mask builder and position-id strategy together.
|
||||
|
||||
Both ``Pipeline`` and ``TokenizeTransform`` need the same triple;
|
||||
centralising the construction avoids drift (e.g. one path forgetting
|
||||
to create the position-id strategy).
|
||||
"""
|
||||
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
|
||||
mask_builder = MaskBuilderFactory.create("sectioned")
|
||||
position_strategy = PositionIdStrategyFactory.create(
|
||||
config.output.position_ids_mode
|
||||
)
|
||||
return tokenizer, mask_builder, position_strategy
|
||||
|
||||
|
||||
def primary_ids(result: dict) -> List[int]:
|
||||
"""Return the first flat int-list value in *result*.
|
||||
|
||||
Used for token counting and position-id generation when the
|
||||
primary key name is not known (DPO uses ``chosen``, GRPO uses
|
||||
``prompts``, SFT uses ``sequence``).
|
||||
"""
|
||||
for val in result.values():
|
||||
if isinstance(val, list) and val and isinstance(val[0], int):
|
||||
return val
|
||||
return []
|
||||
|
||||
|
||||
def infer_dtype(ids: List) -> torch.dtype:
|
||||
"""Float values become float32, everything else int32."""
|
||||
if ids and isinstance(ids[0], float):
|
||||
return torch.float32
|
||||
return torch.int32
|
||||
|
||||
|
||||
def iter_raw_records(
|
||||
records: List[dict],
|
||||
mask_builder,
|
||||
config: PipelineConfig,
|
||||
tokenizer,
|
||||
) -> Iterator[dict]:
|
||||
"""Yield mask-builder output dicts for each record, skipping failures.
|
||||
|
||||
Drops ``domain`` from the result (callers that need it should read
|
||||
it before calling this). Each yielded dict maps a key
|
||||
(``sequence``, ``chosen``, ``responses``…) to either a flat
|
||||
``List[int]`` or a nested ``List[List[int]]`` (GRPO responses/masks).
|
||||
"""
|
||||
for item in records:
|
||||
result = mask_builder.build(item, config, tokenizer)
|
||||
if result is None:
|
||||
continue
|
||||
result.pop("domain", None)
|
||||
if not primary_ids(result):
|
||||
continue
|
||||
yield result
|
||||
|
||||
|
||||
def to_per_record_tensors(
|
||||
raw: Dict[str, list],
|
||||
) -> Dict[str, List[torch.Tensor]]:
|
||||
"""Convert an accumulated ``{key: [per-record ids]}`` dict to tensors.
|
||||
|
||||
Handles three shapes transparently:
|
||||
|
||||
- ``List[int]`` per record (``sequence``, ``chosen``…) → one tensor per record.
|
||||
- ``List[List[int]]`` per record (GRPO ``responses``/``masks``) → one
|
||||
``List[Tensor]`` per record (nested), preserving the per-response
|
||||
boundary so downstream code can index responses individually.
|
||||
- ``List[int]`` for the whole shard (pre-packed keys) → single tensor.
|
||||
|
||||
The detection mirrors the previous inline logic in
|
||||
``Pipeline._flush`` and ``TokenizeTransform.apply``.
|
||||
"""
|
||||
tensors: Dict[str, List[torch.Tensor]] = {}
|
||||
for key, ids_list in raw.items():
|
||||
if ids_list and isinstance(ids_list[0], list):
|
||||
tensors[key] = [
|
||||
[torch.tensor(sub, dtype=infer_dtype(sub)) for sub in ids]
|
||||
if ids and isinstance(ids[0], list)
|
||||
else torch.tensor(ids, dtype=infer_dtype(ids))
|
||||
for ids in ids_list
|
||||
]
|
||||
else:
|
||||
tensors[key] = [
|
||||
torch.tensor(list(chain.from_iterable(ids_list)), dtype=torch.int32)
|
||||
]
|
||||
return tensors
|
||||
|
||||
|
||||
def build_position_ids(
|
||||
sequences: List[List[int]],
|
||||
strategy,
|
||||
) -> Optional[List[int]]:
|
||||
"""Generate position ids for *sequences* using *strategy*.
|
||||
|
||||
Returns ``None`` when the strategy produces no ids (e.g. ``none``
|
||||
mode), so callers can skip attaching the key instead of storing
|
||||
an empty list.
|
||||
"""
|
||||
pos_ids = strategy.generate(sequences)
|
||||
return pos_ids or None
|
||||
Reference in New Issue
Block a user