ViperEkura
31c22dc043
- 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)
93 lines
3.3 KiB
Python
93 lines
3.3 KiB
Python
"""Tokenization transform for JSONL record streams.
|
|
|
|
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
|
|
from pathlib import Path
|
|
from typing import Dict, List
|
|
|
|
import torch
|
|
|
|
from astrai.config.preprocess_config import PipelineConfig
|
|
from astrai.preprocessing.core import (
|
|
build_position_ids,
|
|
build_preprocessing_components,
|
|
iter_raw_records,
|
|
to_per_record_tensors,
|
|
)
|
|
|
|
|
|
class TokenizeTransform:
|
|
"""Tokenize raw JSONL record dicts into per-key tensor lists.
|
|
|
|
Owns the three preprocessing concerns that were previously inlined in
|
|
``JsonlStore``: tokenization, loss-mask construction and position-id
|
|
generation. Constructing it loads the tokenizer, so it is intentionally
|
|
cheap to pass around once built.
|
|
|
|
Args:
|
|
config: Pipeline config describing sections / masks / position mode.
|
|
tokenizer_path: Path passed to ``AutoTokenizer.from_pretrained``.
|
|
"""
|
|
|
|
def __init__(self, config: PipelineConfig, tokenizer_path: str):
|
|
self.config = config
|
|
self.tokenizer, self.mask_builder, self.position_strategy = (
|
|
build_preprocessing_components(config, tokenizer_path)
|
|
)
|
|
|
|
@classmethod
|
|
def from_config_file(cls, config_path: str) -> "TokenizeTransform":
|
|
"""Build from a ``dataset_config.json`` file path.
|
|
|
|
The config file follows :class:`PipelineConfig` schema with an
|
|
extra ``tokenizer_path`` field. When omitted, the config's
|
|
parent directory is used as the tokenizer path.
|
|
"""
|
|
root = Path(config_path).parent
|
|
with open(config_path, "r", encoding="utf-8") as f:
|
|
raw_config = json.load(f)
|
|
tokenizer_path = raw_config.pop("tokenizer_path", None) or str(root)
|
|
config = PipelineConfig.from_dict(raw_config)
|
|
return cls(config, tokenizer_path)
|
|
|
|
def apply(self, records: List[dict]) -> Dict[str, list]:
|
|
"""Tokenize a list of raw record dicts.
|
|
|
|
Returns a dict mapping key (``sequence``, ``chosen``, ``responses``,
|
|
…) to a list of per-record tensors (or nested tensor lists for
|
|
multi-response keys such as GRPO ``responses``).
|
|
"""
|
|
raw: Dict[str, list] = {}
|
|
doc_sequences: List[List[int]] = []
|
|
|
|
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():
|
|
raw.setdefault(key, []).append(ids)
|
|
|
|
tensors = to_per_record_tensors(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)]
|
|
|
|
return tensors
|