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:
2026-07-19 12:27:56 +08:00
parent 17127f8b3c
commit 31c22dc043
6 changed files with 297 additions and 146 deletions
+29 -40
View File
@@ -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