diff --git a/astrai/preprocessing/__init__.py b/astrai/preprocessing/__init__.py index 31bc5f3..3362acc 100644 --- a/astrai/preprocessing/__init__.py +++ b/astrai/preprocessing/__init__.py @@ -8,6 +8,7 @@ from astrai.preprocessing.builder import ( from astrai.preprocessing.packing import ( PackingStrategy, PackingStrategyFactory, + plan_bfd, ) from astrai.preprocessing.pipeline import Pipeline, filter_by_length from astrai.preprocessing.position_id import ( @@ -35,4 +36,5 @@ __all__ = [ "StoreWriterFactory", "TokenizeTransform", "filter_by_length", + "plan_bfd", ] diff --git a/astrai/preprocessing/core.py b/astrai/preprocessing/core.py new file mode 100644 index 0000000..50cf118 --- /dev/null +++ b/astrai/preprocessing/core.py @@ -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 diff --git a/astrai/preprocessing/packing.py b/astrai/preprocessing/packing.py index 05e2621..186d41e 100644 --- a/astrai/preprocessing/packing.py +++ b/astrai/preprocessing/packing.py @@ -19,6 +19,43 @@ def _truncate(seq: List[int], max_len: int, mode: str) -> List[int]: return seq[:max_len] +def plan_bfd( + sequences: List[List[int]], max_packed_len: int, truncation_mode: str = "keep_start" +) -> List[List[int]]: + """Best-Fit Decreasing bin packing of *sequences* into bins. + + Returns a list of bins, each bin a list of original indices into + *sequences*. Bin capacities are respected on the *truncated* + length of each sequence (so a sequence longer than + *max_packed_len* counts at *max_packed_len*). + + Pure index-based so callers can apply the same plan to any + aligned key (``loss_mask``, ``position_ids``…). + """ + n = len(sequences) + order = sorted(range(n), key=lambda i: len(sequences[i]), reverse=True) + bins: List[List[int]] = [] + bin_lengths: List[int] = [] + + for orig_idx in order: + seq_len = len(_truncate(sequences[orig_idx], max_packed_len, truncation_mode)) + best_bin = None + best_remain = max_packed_len + 1 + for i, bl in enumerate(bin_lengths): + remain = max_packed_len - bl + if seq_len <= remain < best_remain: + best_remain = remain + best_bin = i + if best_bin is not None: + bins[best_bin].append(orig_idx) + bin_lengths[best_bin] += seq_len + else: + bins.append([orig_idx]) + bin_lengths.append(seq_len) + + return bins + + class PackingStrategy(ABC): """Reorder and truncate sequences within a shard.""" @@ -70,7 +107,7 @@ class BFDPacking(PackingStrategy): sequences = keys.get("sequence", []) if not sequences: return keys - bins = self._plan(sequences, max_packed_len, truncation_mode) + bins = plan_bfd(sequences, max_packed_len, truncation_mode) packed: Dict[str, List[List[int]]] = {} for k, vals in keys.items(): @@ -91,35 +128,6 @@ class BFDPacking(PackingStrategy): result.extend(vals[i]) return result - @staticmethod - def _plan( - sequences: List[List[int]], max_packed_len: int, truncation_mode: str - ) -> List[List[int]]: - n = len(sequences) - order = sorted(range(n), key=lambda i: len(sequences[i]), reverse=True) - bins: List[List[int]] = [] - bin_lengths: List[int] = [] - - for orig_idx in order: - seq_len = len( - _truncate(sequences[orig_idx], max_packed_len, truncation_mode) - ) - best_bin = None - best_remain = max_packed_len + 1 - for i, bl in enumerate(bin_lengths): - remain = max_packed_len - bl - if seq_len <= remain < best_remain: - best_remain = remain - best_bin = i - if best_bin is not None: - bins[best_bin].append(orig_idx) - bin_lengths[best_bin] += seq_len - else: - bins.append([orig_idx]) - bin_lengths.append(seq_len) - - return bins - @PackingStrategyFactory.register("bfd_split") class BFDSplitPacking(BFDPacking): diff --git a/astrai/preprocessing/pipeline.py b/astrai/preprocessing/pipeline.py index fca8c21..cfcf8e1 100644 --- a/astrai/preprocessing/pipeline.py +++ b/astrai/preprocessing/pipeline.py @@ -4,6 +4,10 @@ Composes a :class:`BaseMaskBuilder` (selected by ``input.type``) with sharding and flush to ``.h5`` / ``.bin`` storage. Packing, position-id generation and storage writing are each delegated to pluggable strategies, dispatched by configuration keys. + +Record iteration, mask building, primary-id extraction and per-key +accumulation are shared with :class:`TokenizeTransform` via the +:mod:`astrai.preprocessing.core` helpers. """ import json @@ -17,11 +21,13 @@ import torch import tqdm from astrai.config.preprocess_config import PipelineConfig -from astrai.preprocessing.builder import MaskBuilderFactory +from astrai.preprocessing.core import ( + build_preprocessing_components, + iter_raw_records, + primary_ids, +) from astrai.preprocessing.packing import PackingStrategyFactory -from astrai.preprocessing.position_id import PositionIdStrategyFactory from astrai.preprocessing.writer import StoreWriterFactory -from astrai.tokenize import AutoTokenizer logger = logging.getLogger(__name__) @@ -64,20 +70,18 @@ class Pipeline: self.output_dir = output_dir self.tokenizer_path = tokenizer_path - self.mask_builder = MaskBuilderFactory.create("sectioned") + self.tokenizer, self.mask_builder, self._position_id = ( + build_preprocessing_components(config, tokenizer_path) + ) self._packer = PackingStrategyFactory.create( config.preprocessing.packing_strategy ) - self._position_id = PositionIdStrategyFactory.create( - config.output.position_ids_mode - ) self._writer = StoreWriterFactory.create(config.output.storage_format) def transform(self, item: dict) -> Optional[dict]: - return self.mask_builder.build(item, self.config, self._tokenizer) + return self.mask_builder.build(item, self.config, self.tokenizer) def run(self): - self._tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_path) domains: dict = defaultdict(lambda: defaultdict(list)) total_tokens = 0 shard_idx: dict[str, int] = defaultdict(int) @@ -102,14 +106,7 @@ class Pipeline: continue 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 - + ids = primary_ids(result) if not ids: continue @@ -129,15 +126,6 @@ class Pipeline: if total_tokens > 0: self._flush(domains, shard_idx) - @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): """Pad previously-accumulated keys that are missing from *result*.""" @@ -170,39 +158,12 @@ class Pipeline: original_sequences = keys.get("sequence", []) mode = self.config.output.position_ids_mode - if mode == "doc_reset" and original_sequences: - keys["position_ids"] = [list(range(len(s))) for s in original_sequences] - + keys = self._inject_doc_reset_position_ids(keys, mode, original_sequences) keys = self._packer.apply(dict(keys), pp.max_packed_len, pp.truncation_mode) - - tensors: Dict[str, List[torch.Tensor]] = {} - for key, ids_list in keys.items(): - dt = _STR_TO_DTYPE.get( - self.config.output.dtype.get(key, "int32"), torch.int32 - ) - # GRPO multi-response keys store List[List[int]] per record - # (responses/masks). Rewards store List[float] per record. - # Both produce List[Tensor] (one tensor per record), but - # responses need inner flattening while rewards do not. - if ids_list and isinstance(ids_list[0], list): - tensors[key] = [ - torch.tensor( - list(chain.from_iterable(ids)) - if ids and isinstance(ids[0], list) - else ids, - dtype=dt, - ) - for ids in ids_list - ] - else: - tensors[key] = [ - torch.tensor(list(chain.from_iterable(ids_list)), dtype=dt) - ] - - if mode == "continuous" and original_sequences: - pos_ids = self._position_id.generate(keys.get("sequence", [])) - if pos_ids: - tensors["position_ids"] = [torch.tensor(pos_ids, dtype=torch.int32)] + tensors = self._to_tensors(keys) + tensors = self._inject_continuous_position_ids( + tensors, mode, keys.get("sequence", []) + ) self._writer.save(self.output_dir, domain, idx, tensors) shard_idx[domain] = idx + 1 @@ -212,3 +173,76 @@ class Pipeline: f" saved {domain}/shard_{idx:04d} " f"({tensors[first_key][0].numel():,} tokens)" ) + + def _inject_doc_reset_position_ids( + self, + keys: Dict[str, list], + mode: str, + original_sequences: List[List[int]], + ) -> Dict[str, list]: + """Attach per-document position_ids before packing (``doc_reset``). + + ``doc_reset`` position ids must enter the packer so that each + packed bin concatenates the per-doc ranges in bin order. The + per-record structure ``[range(len(s)) for s in seqs]`` is required + by the packer (it concatenates per-record lists per bin); the + ``PositionIdStrategy.generate`` flattens, so it cannot be used + directly here — it is only consulted for the ``continuous`` + post-packing path. + """ + if mode != "doc_reset" or not original_sequences: + return keys + keys["position_ids"] = [list(range(len(s))) for s in original_sequences] + return keys + + def _inject_continuous_position_ids( + self, + tensors: Dict[str, List[torch.Tensor]], + mode: str, + packed_sequences: List[List[int]], + ) -> Dict[str, List[torch.Tensor]]: + """Attach a single continuous position_ids tensor after packing. + + ``continuous`` mode spans the whole shard (post-packing), so it + cannot participate in bin packing — it is computed from the + packed sequences and appended directly to the tensor dict. + """ + if mode != "continuous" or not packed_sequences: + return tensors + pos_ids = self._position_id.generate(packed_sequences) + if pos_ids: + tensors["position_ids"] = [torch.tensor(pos_ids, dtype=torch.int32)] + return tensors + + def _to_tensors(self, keys: Dict[str, list]) -> Dict[str, List[torch.Tensor]]: + """Convert packed per-key id lists to tensors. + + Honours ``config.output.dtype`` overrides per key; falls back to + ``int32``. Handles three shapes (see + :func:`astrai.preprocessing.core.to_per_record_tensors` for the + equivalent online-path helper): + - ``List[int]`` per record → one tensor per record. + - ``List[List[int]]`` per record (GRPO responses/masks) → one tensor + per record, inner lists flattened. + - ``List[int]`` for the whole shard (pre-packed keys) → single tensor. + """ + tensors: Dict[str, List[torch.Tensor]] = {} + for key, ids_list in keys.items(): + dt = _STR_TO_DTYPE.get( + self.config.output.dtype.get(key, "int32"), torch.int32 + ) + if ids_list and isinstance(ids_list[0], list): + tensors[key] = [ + torch.tensor( + list(chain.from_iterable(ids)) + if ids and isinstance(ids[0], list) + else ids, + dtype=dt, + ) + for ids in ids_list + ] + else: + tensors[key] = [ + torch.tensor(list(chain.from_iterable(ids_list)), dtype=dt) + ] + return tensors diff --git a/astrai/preprocessing/transform.py b/astrai/preprocessing/transform.py index c95ab8c..f56ec4a 100644 --- a/astrai/preprocessing/transform.py +++ b/astrai/preprocessing/transform.py @@ -4,6 +4,11 @@ Bridges the Reader layer (``JsonlStore`` reads raw JSON records) and the Dataset layer (expects per-record tensors). Holds the tokenizer, mask-builder and position-id strategy together so that I/O code stays free of model dependencies. + +The record-processing core (mask building, primary-id extraction, +per-key tensorisation, position-id generation) is shared with +:class:`astrai.preprocessing.pipeline.Pipeline` via the +:mod:`astrai.preprocessing.core` helpers. """ import json @@ -13,9 +18,12 @@ from typing import Dict, List 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 +from astrai.preprocessing.core import ( + build_position_ids, + build_preprocessing_components, + iter_raw_records, + to_per_record_tensors, +) class TokenizeTransform: @@ -33,10 +41,8 @@ class TokenizeTransform: def __init__(self, config: PipelineConfig, tokenizer_path: str): self.config = config - self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) - self.mask_builder = MaskBuilderFactory.create("sectioned") - self.position_strategy = PositionIdStrategyFactory.create( - config.output.position_ids_mode + self.tokenizer, self.mask_builder, self.position_strategy = ( + build_preprocessing_components(config, tokenizer_path) ) @classmethod @@ -64,40 +70,23 @@ class TokenizeTransform: raw: Dict[str, list] = {} doc_sequences: List[List[int]] = [] - for item in records: - result = self.mask_builder.build(item, self.config, self.tokenizer) - if result is None: - continue - result.pop("domain", None) - primary_ids = self._primary_ids(result) - if not primary_ids: - continue - doc_sequences.append(primary_ids) + for result in iter_raw_records( + records, self.mask_builder, self.config, self.tokenizer + ): + primary = None + for val in result.values(): + if isinstance(val, list) and val and isinstance(val[0], int): + primary = val + break + if primary is not None: + doc_sequences.append(primary) for key, ids in result.items(): - if key not in raw: - raw[key] = [] - if ids and isinstance(ids[0], list): - raw[key].append( - [torch.tensor(sub, dtype=self._infer_dtype(sub)) for sub in ids] - ) - else: - raw[key].append(torch.tensor(ids, dtype=self._infer_dtype(ids))) + raw.setdefault(key, []).append(ids) - pos_ids = self.position_strategy.generate(doc_sequences) - if pos_ids: - raw["position_ids"] = [torch.tensor(pos_ids, dtype=torch.int32)] + tensors = to_per_record_tensors(raw) - return raw + pos_ids = build_position_ids(doc_sequences, self.position_strategy) + if pos_ids is not None: + tensors["position_ids"] = [torch.tensor(pos_ids, dtype=torch.int32)] - @staticmethod - def _primary_ids(result: dict) -> List[int]: - for val in result.values(): - if isinstance(val, list) and val and isinstance(val[0], int): - return val - return [] - - @staticmethod - def _infer_dtype(ids: List) -> torch.dtype: - if ids and isinstance(ids[0], float): - return torch.float32 - return torch.int32 + return tensors diff --git a/scripts/eval/evaluate_ifd.py b/scripts/eval/evaluate_ifd.py index bc0a164..366b0f0 100644 --- a/scripts/eval/evaluate_ifd.py +++ b/scripts/eval/evaluate_ifd.py @@ -26,28 +26,22 @@ import torch.nn.functional as F import tqdm from astrai.model import AutoModel +from astrai.preprocessing.packing import plan_bfd from astrai.tokenize import AutoTokenizer def _pack_bins(pairs, max_len): - """BFD bin packing: pack (c+r) into bins of max total length.""" - indexed = sorted(enumerate(pairs), key=lambda x: -(len(x[1][0]) + len(x[1][1]))) - bins = [] - lengths = [] - for orig_idx, (c, r) in indexed: - size = len(c) + len(r) - best_bin = -1 - for bi, rem in enumerate(lengths): - if rem >= size: - if best_bin < 0 or rem < lengths[best_bin]: - best_bin = bi - if best_bin >= 0: - bins[best_bin].append((orig_idx, c, r)) - lengths[best_bin] -= size - else: - bins.append([(orig_idx, c, r)]) - lengths.append(max_len - size) - return bins + """BFD bin packing: pack (c+r) into bins of max total length. + + Reuses :func:`plan_bfd` so the BFD heuristic stays single-sourced. + """ + # Treat each pair as a single sequence of length len(c)+len(r) for + # planning purposes; plan_bfd works on pure lengths. + fake_sequences = [[0] * (len(c) + len(r)) for c, r in pairs] + plan = plan_bfd(fake_sequences, max_len) + return [ + [(i, pairs[i][0], pairs[i][1]) for i in bin_indices] for bin_indices in plan + ] def _resolve_sentinel_ids(tokenizer, sentinel_text):