ViperEkura
f163520fff
- move JSONL transform auto-creation from JsonlStore.load to DatasetFactory.load via _build_jsonl_transform helper - remove TokenizeTransform and PipelineConfig imports from storage module - JsonlStore.load now requires explicit transform= for eager mode - DatasetFactory.load remains the public API with identical convenience behavior
539 lines
19 KiB
Python
539 lines
19 KiB
Python
"""Dataset implementations for training.
|
|
|
|
Composition over inheritance — every dataset is a thin wrapper that
|
|
binds a :class:`Store` to a particular train-type's key mapping. All
|
|
sample-id → token/record indexing lives on the Store; datasets never
|
|
know about window/stride math or segment layouts.
|
|
|
|
Class hierarchy:
|
|
|
|
BaseDataset (ABC) — holds a Store, exposes __len__/keys,
|
|
overrides __getitem__
|
|
├── SEQDataset — next-token prediction (stream)
|
|
├── SFTDataset — loss-mask + position_ids (stream)
|
|
├── DPODataset — chosen/rejected pairs (record)
|
|
└── GRPODataset — prompt + response group (record)
|
|
|
|
``DatasetFactory.load(train_type, load_path, window_size, stride, …)``
|
|
builds the Store (auto-detecting format) before constructing the
|
|
matching dataset. Passing ``store=`` skips Store construction.
|
|
|
|
When a record dataset (DPO) reads from raw JSONL, a *processor*
|
|
function (pure ``record -> Dict[str, Tensor]``) is forwarded to
|
|
:class:`JsonlStore` so tokenisation happens on the fly.
|
|
"""
|
|
|
|
from abc import ABC, abstractmethod
|
|
from functools import partial
|
|
from pathlib import Path
|
|
from typing import Callable, Dict, List, Optional
|
|
|
|
import torch
|
|
from torch import Tensor
|
|
from torch.utils.data import Dataset
|
|
|
|
from astrai.config.preprocess_config import PipelineConfig
|
|
from astrai.dataset.storage import (
|
|
Store,
|
|
StoreFactory,
|
|
detect_format,
|
|
)
|
|
from astrai.factory import BaseFactory
|
|
from astrai.preprocessing.transform import TokenizeTransform
|
|
from astrai.tokenize import AutoTokenizer
|
|
|
|
_DEFAULT_MESSAGES_CONFIG = {
|
|
"version": 1,
|
|
"input": {"sections": [{"field": "messages", "action": "$role", "template": True}]},
|
|
"mask": {"system": "mask", "user": "mask", "assistant": "train"},
|
|
"mask_default": "mask",
|
|
"output": {"position_ids_mode": "doc_reset"},
|
|
}
|
|
|
|
|
|
def _build_jsonl_transform(
|
|
path: str, tokenizer_path: Optional[str] = None
|
|
) -> Optional["TokenizeTransform"]:
|
|
"""Auto-build a TokenizeTransform for JSONL eager loading.
|
|
|
|
Reads ``dataset_config.json`` from the data dir if present, or
|
|
falls back to the built-in chatml SFT config when *tokenizer_path*
|
|
is provided.
|
|
"""
|
|
root = Path(path)
|
|
config_path = root / "dataset_config.json" if root.is_dir() else None
|
|
if config_path is not None and config_path.exists():
|
|
return TokenizeTransform.from_config_file(str(config_path))
|
|
if tokenizer_path:
|
|
config = PipelineConfig.from_dict(_DEFAULT_MESSAGES_CONFIG)
|
|
return TokenizeTransform(config, tokenizer_path)
|
|
return None
|
|
|
|
|
|
def dpo_tokenize(
|
|
record: dict,
|
|
tokenizer,
|
|
max_len: int = 2048,
|
|
) -> Optional[dict]:
|
|
"""Tokenize one DPO record into chosen/rejected + masks.
|
|
|
|
Applies the tokenizer's chat template so token sequences match the
|
|
SFT checkpoint's format. Prompt is rendered with
|
|
``add_generation_prompt=True``; chosen/rejected are appended as a
|
|
single assistant turn.
|
|
|
|
Accepts:
|
|
|
|
- Flat: ``{"prompt": str, "chosen": str, "rejected": str}``
|
|
- Conv: ``{"prompt": [{role, content}, ...], "chosen": [...], ...}``
|
|
- Legacy: ``{"input": str, "chosen": str, "rejected": str}``
|
|
|
|
No packing, no ``position_ids`` — DPO sequences are independent.
|
|
"""
|
|
prompt = record.get("prompt") or record.get("input")
|
|
chosen = record.get("chosen")
|
|
rejected = record.get("rejected")
|
|
if prompt is None or chosen is None or rejected is None:
|
|
return None
|
|
|
|
prompt_messages = _to_messages(prompt)
|
|
chosen_text = _extract_text(chosen)
|
|
rejected_text = _extract_text(rejected)
|
|
if chosen_text is None or rejected_text is None:
|
|
return None
|
|
chosen_messages = prompt_messages + [{"role": "assistant", "content": chosen_text}]
|
|
rejected_messages = prompt_messages + [
|
|
{"role": "assistant", "content": rejected_text}
|
|
]
|
|
|
|
prompt_ids = tokenizer.apply_chat_template(
|
|
prompt_messages, tokenize=True, add_generation_prompt=True
|
|
)
|
|
ch_ids = tokenizer.apply_chat_template(
|
|
chosen_messages, tokenize=True, add_generation_prompt=False
|
|
)
|
|
re_ids = tokenizer.apply_chat_template(
|
|
rejected_messages, tokenize=True, add_generation_prompt=False
|
|
)
|
|
|
|
full_ch = ch_ids[:max_len]
|
|
full_re = re_ids[:max_len]
|
|
|
|
prompt_len = min(len(prompt_ids), max_len)
|
|
ch_mask = [0] * prompt_len + [1] * max(0, len(full_ch) - prompt_len)
|
|
ch_mask = ch_mask[:max_len]
|
|
re_mask = [0] * prompt_len + [1] * max(0, len(full_re) - prompt_len)
|
|
re_mask = re_mask[:max_len]
|
|
|
|
return {
|
|
"chosen": full_ch,
|
|
"rejected": full_re,
|
|
"chosen_mask": ch_mask,
|
|
"rejected_mask": re_mask,
|
|
}
|
|
|
|
|
|
def _to_messages(value) -> list:
|
|
"""Accept str or conversation list; return message list."""
|
|
if isinstance(value, str):
|
|
return [{"role": "user", "content": value}]
|
|
if isinstance(value, list):
|
|
return value
|
|
return [{"role": "user", "content": str(value)}]
|
|
|
|
|
|
def _extract_text(value) -> Optional[str]:
|
|
"""Accept str or conversation list; return plain text."""
|
|
if value is None:
|
|
return None
|
|
if isinstance(value, str):
|
|
return value
|
|
if isinstance(value, list):
|
|
return "".join(m.get("content", "") for m in value if isinstance(m, dict))
|
|
return None
|
|
|
|
|
|
def dpo_processor(
|
|
record: dict,
|
|
tokenizer,
|
|
max_len: int = 2048,
|
|
) -> Dict[str, Tensor]:
|
|
"""DPO processor: wraps :func:`dpo_tokenize` and returns tensors."""
|
|
result = dpo_tokenize(record, tokenizer, max_len=max_len)
|
|
if result is None:
|
|
raise ValueError(f"Malformed DPO record: {list(record.keys())}")
|
|
return {
|
|
"chosen": torch.tensor(result["chosen"], dtype=torch.int32),
|
|
"rejected": torch.tensor(result["rejected"], dtype=torch.int32),
|
|
"chosen_mask": torch.tensor(result["chosen_mask"], dtype=torch.bool),
|
|
"rejected_mask": torch.tensor(result["rejected_mask"], dtype=torch.bool),
|
|
}
|
|
|
|
|
|
def dpo_collate_fn(batch: List[Dict[str, Tensor]]) -> Dict[str, Tensor]:
|
|
"""Collate variable-length DPO samples into padded 2-D tensors.
|
|
|
|
Input: list of dicts, each with:
|
|
- chosen: [C_i]
|
|
- rejected: [R_i]
|
|
- chosen_mask: [C_i]
|
|
- rejected_mask: [R_i]
|
|
|
|
Output (padded to the max length across chosen/rejected within the batch):
|
|
- chosen: [B, S_max]
|
|
- rejected: [B, S_max]
|
|
- chosen_mask: [B, S_max]
|
|
- rejected_mask: [B, S_max]
|
|
"""
|
|
B = len(batch)
|
|
S_max = max(b["chosen"].size(0) for b in batch)
|
|
S_max = max(S_max, max(b["rejected"].size(0) for b in batch))
|
|
|
|
chosen = torch.zeros(B, S_max, dtype=torch.long)
|
|
rejected = torch.zeros(B, S_max, dtype=torch.long)
|
|
chosen_mask = torch.zeros(B, S_max, dtype=torch.bool)
|
|
rejected_mask = torch.zeros(B, S_max, dtype=torch.bool)
|
|
|
|
for i, b in enumerate(batch):
|
|
c_len = b["chosen"].size(0)
|
|
r_len = b["rejected"].size(0)
|
|
chosen[i, :c_len] = b["chosen"]
|
|
rejected[i, :r_len] = b["rejected"]
|
|
chosen_mask[i, :c_len] = b["chosen_mask"]
|
|
rejected_mask[i, :r_len] = b["rejected_mask"]
|
|
|
|
return {
|
|
"chosen": chosen,
|
|
"rejected": rejected,
|
|
"chosen_mask": chosen_mask,
|
|
"rejected_mask": rejected_mask,
|
|
}
|
|
|
|
|
|
def grpo_collate_fn(batch: List[Dict[str, Tensor]]) -> Dict[str, Tensor]:
|
|
"""Collate variable-length GRPO samples into padded 3-D tensors.
|
|
|
|
Input: list of dicts, each with:
|
|
- prompts: [P_i]
|
|
- responses: list of G tensors, each [R_ij]
|
|
- masks: list of G tensors, each [R_ij]
|
|
- rewards: [G]
|
|
|
|
Output:
|
|
- prompts: [B, P_max], left-padded
|
|
- prompt_mask: [B, P_max]
|
|
- responses: [B, G, R_max]
|
|
- masks: [B, G, R_max]
|
|
- rewards: [B, G]
|
|
"""
|
|
B = len(batch)
|
|
G = len(batch[0]["responses"])
|
|
P_max = max(b["prompts"].size(0) for b in batch)
|
|
R_max = max(r.size(0) for b in batch for r in b["responses"])
|
|
|
|
prompts = torch.zeros(B, P_max, dtype=torch.long)
|
|
prompt_mask = torch.zeros(B, P_max, dtype=torch.bool)
|
|
responses = torch.zeros(B, G, R_max, dtype=torch.long)
|
|
masks = torch.zeros(B, G, R_max, dtype=torch.bool)
|
|
rewards = torch.zeros(B, G, dtype=torch.float32)
|
|
|
|
for i, b in enumerate(batch):
|
|
p_len = b["prompts"].size(0)
|
|
prompts[i, -p_len:] = b["prompts"]
|
|
prompt_mask[i, -p_len:] = True
|
|
rewards[i, : b["rewards"].size(0)] = b["rewards"]
|
|
for g in range(min(G, len(b["responses"]))):
|
|
r_len = b["responses"][g].size(0)
|
|
responses[i, g, :r_len] = b["responses"][g]
|
|
if g < len(b["masks"]):
|
|
masks[i, g, :r_len] = b["masks"][g]
|
|
|
|
return {
|
|
"prompts": prompts,
|
|
"prompt_mask": prompt_mask,
|
|
"responses": responses,
|
|
"masks": masks,
|
|
"rewards": rewards,
|
|
}
|
|
|
|
|
|
def validate_keys(store: Store, required: List[str]) -> None:
|
|
"""Raise ``KeyError`` if *store* is missing any *required* key."""
|
|
if not required:
|
|
return
|
|
actual = set(store.keys)
|
|
missing = [k for k in required if k not in actual]
|
|
if missing:
|
|
raise KeyError(
|
|
f"Store at {getattr(store, '_load_path', '?')} is missing required "
|
|
f"keys {missing}; available keys are {sorted(actual)}."
|
|
)
|
|
|
|
|
|
class BaseDataset(Dataset, ABC):
|
|
"""Abstract base class for dataset types.
|
|
|
|
Holds a :class:`Store`. All sample-id indexing is delegated to the
|
|
store — this class exposes ``__len__`` as ``len(store)`` and the
|
|
``keys`` property as ``store.keys``. Subclasses implement
|
|
``__getitem__`` with the train-type-specific key mapping and any
|
|
training-only index arithmetic (e.g. the next-token ``+1`` shift).
|
|
"""
|
|
|
|
required_keys: List[str] = []
|
|
|
|
def __init__(self, store: Store):
|
|
super().__init__()
|
|
self.store: Store = store
|
|
validate_keys(store, self.required_keys)
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.store)
|
|
|
|
@property
|
|
def keys(self) -> List[str]:
|
|
return self.store.keys
|
|
|
|
@property
|
|
def token_count(self) -> int:
|
|
return self.store.token_count
|
|
|
|
@abstractmethod
|
|
def __getitem__(self, index: int) -> Dict[str, Tensor]:
|
|
raise NotImplementedError
|
|
|
|
|
|
class DatasetFactory(BaseFactory["BaseDataset"]):
|
|
"""Factory for creating dataset instances by train-type.
|
|
|
|
Use :meth:`DatasetFactory.register("custom")` to register new
|
|
dataset classes; they must inherit from :class:`BaseDataset`.
|
|
"""
|
|
|
|
@classmethod
|
|
def load(
|
|
cls,
|
|
train_type: str,
|
|
load_path: Optional[str] = None,
|
|
window_size: int = 0,
|
|
stride: Optional[int] = None,
|
|
storage_type: Optional[str] = None,
|
|
tokenizer_path: Optional[str] = None,
|
|
max_len: int = 2048,
|
|
store: Optional[Store] = None,
|
|
**kwargs,
|
|
) -> "BaseDataset":
|
|
"""Create and load a dataset in one step.
|
|
|
|
Two entry points:
|
|
|
|
- **store given**: bind it directly — the caller fully controls
|
|
Store construction and processor setup. *load_path*,
|
|
*storage_type*, *tokenizer_path*, *window_size*, *stride* are
|
|
ignored.
|
|
- **store is None**: build a Store from *load_path*, auto-detecting
|
|
format and constructing a processor when *tokenizer_path* is
|
|
given for a record dataset on JSONL.
|
|
|
|
Args:
|
|
train_type: Registered dataset name ("seq", "sft", "dpo",
|
|
"grpo", …).
|
|
load_path: Path to the data file or directory (ignored if
|
|
*store* is given).
|
|
window_size: Stream window length — only meaningful for
|
|
stream datasets (SEQ/SFT). Record datasets ignore it.
|
|
stride: Stride between consecutive stream samples
|
|
(default: same as *window_size*).
|
|
storage_type: Storage backend ("bin", "jsonl") or
|
|
None for auto-detection.
|
|
tokenizer_path: Path to tokenizer for lazy JSONL
|
|
tokenisation (record datasets only).
|
|
max_len: Max sequence length forwarded to processors.
|
|
store: Pre-built, already-loaded Store instance.
|
|
**kwargs: Extra arguments forwarded to ``store.load()``.
|
|
|
|
Returns:
|
|
Loaded dataset instance.
|
|
"""
|
|
if store is not None:
|
|
return cls.create(train_type, store=store)
|
|
|
|
if load_path is None:
|
|
raise ValueError("Either load_path or store must be provided")
|
|
|
|
if storage_type is None:
|
|
storage_type = detect_format(load_path)
|
|
|
|
if stride is None:
|
|
stride = window_size
|
|
|
|
processor = cls._maybe_build_processor(
|
|
train_type, storage_type, tokenizer_path, max_len
|
|
)
|
|
|
|
store_window = cls._store_window_for(train_type, window_size)
|
|
store = StoreFactory.create(
|
|
storage_type,
|
|
window_size=store_window,
|
|
stride=stride if stride else store_window,
|
|
)
|
|
if processor is not None:
|
|
store.load(load_path, processor=processor, **kwargs)
|
|
elif storage_type == "jsonl":
|
|
transform = _build_jsonl_transform(load_path, tokenizer_path)
|
|
if transform is None:
|
|
raise FileNotFoundError(
|
|
f"JSONL dataset config not found. Expected "
|
|
f"dataset_config.json alongside *.jsonl files, pass "
|
|
f"tokenizer_path= for the built-in messages config, or "
|
|
f"use processor= for lazy on-the-fly tokenisation."
|
|
)
|
|
store.load(load_path, transform=transform, **kwargs)
|
|
else:
|
|
store.load(load_path, **kwargs)
|
|
|
|
return cls.create(train_type, store=store)
|
|
|
|
@staticmethod
|
|
def _store_window_for(train_type: str, window_size: int) -> int:
|
|
"""Stream datasets consume ``window_size``; record datasets ignore it.
|
|
|
|
Record datasets (dpo/grpo) treat each record as an independent
|
|
training unit and never window, so the store is built with
|
|
``window_size=0`` and ``len(store)`` returns the record count.
|
|
"""
|
|
if train_type in ("seq", "sft"):
|
|
return window_size
|
|
return 0
|
|
|
|
@staticmethod
|
|
def _maybe_build_processor(
|
|
train_type: str,
|
|
storage_type: str,
|
|
tokenizer_path: Optional[str],
|
|
max_len: int,
|
|
) -> Optional[Callable[[dict], Dict[str, Tensor]]]:
|
|
"""Build an on-the-fly tokenisation processor if applicable.
|
|
|
|
Only raw JSONL + record datasets (DPO/GRPO) need a processor;
|
|
pre-tokenised backends (bin) and stream datasets (SEQ/SFT)
|
|
return ``None`` so no tokenizer is loaded.
|
|
"""
|
|
if tokenizer_path is None or storage_type != "jsonl":
|
|
return None
|
|
if train_type == "dpo":
|
|
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
|
|
return partial(dpo_processor, tokenizer=tokenizer, max_len=max_len)
|
|
return None
|
|
|
|
|
|
@DatasetFactory.register("seq")
|
|
class SEQDataset(BaseDataset):
|
|
"""Dataset for sequential next-token prediction training.
|
|
|
|
Stream mode: ``store.fetch(begin, end, "sequence")`` returns the
|
|
input window; the +1 shifted call returns the next-token target.
|
|
"""
|
|
|
|
required_keys = ["sequence"]
|
|
|
|
def __getitem__(self, index: int):
|
|
begin, end = self.store.sample_window(index)
|
|
x = self.store.fetch(begin, end, "sequence")
|
|
y = self.store.fetch(begin + 1, end + 1, "sequence")
|
|
return {
|
|
"input_ids": x.to(dtype=torch.long),
|
|
"target_ids": y.to(dtype=torch.long),
|
|
}
|
|
|
|
|
|
@DatasetFactory.register("sft")
|
|
class SFTDataset(BaseDataset):
|
|
"""Dataset for supervised fine-tuning with loss masking.
|
|
|
|
Stream mode: ``sequence``/``loss_mask``/``position_ids`` are sliced
|
|
to the window. ``loss_mask`` and ``target_ids`` use the +1 shifted
|
|
slice so they align with the predicted positions.
|
|
"""
|
|
|
|
required_keys = ["sequence", "loss_mask", "position_ids"]
|
|
|
|
def __getitem__(self, index: int):
|
|
begin, end = self.store.sample_window(index)
|
|
x = self.store.fetch(begin, end, "sequence")
|
|
y = self.store.fetch(begin + 1, end + 1, "sequence")
|
|
position_ids = self.store.fetch(begin, end, "position_ids")
|
|
loss_mask = self.store.fetch(begin + 1, end + 1, "loss_mask")
|
|
return {
|
|
"input_ids": x.to(dtype=torch.long),
|
|
"target_ids": y.to(dtype=torch.long),
|
|
"position_ids": position_ids.to(dtype=torch.long),
|
|
"loss_mask": loss_mask.to(dtype=torch.bool),
|
|
}
|
|
|
|
|
|
@DatasetFactory.register("dpo")
|
|
class DPODataset(BaseDataset):
|
|
"""Record-structured dataset for Direct Preference Optimization.
|
|
|
|
Each sample is one preference pair (chosen + rejected) and is an
|
|
independent training unit — no windowing, stride, or cross-record
|
|
concatenation. This keeps each sequence self-contained so attention
|
|
never leaks across preference pairs.
|
|
|
|
Two loading paths (handled by :class:`DatasetFactory`):
|
|
|
|
- **Pre-tokenized** (bin): ``store.load(path)`` reads per-record
|
|
tensors; ``__getitem__`` returns them directly.
|
|
- **Raw JSONL** (``tokenizer_path=...``): builds a lazy processor
|
|
via :func:`dpo_processor` that tokenises on the fly — no packing,
|
|
no ``position_ids``.
|
|
"""
|
|
|
|
required_keys = ["chosen", "rejected", "chosen_mask", "rejected_mask"]
|
|
|
|
def make_processor(self, tokenizer, max_len: int):
|
|
return partial(dpo_processor, tokenizer=tokenizer, max_len=max_len)
|
|
|
|
def __getitem__(self, index: int) -> Dict[str, Tensor]:
|
|
return {
|
|
"chosen": self.store.fetch_record(index, "chosen").to(dtype=torch.long),
|
|
"rejected": self.store.fetch_record(index, "rejected").to(dtype=torch.long),
|
|
"chosen_mask": self.store.fetch_record(index, "chosen_mask").to(
|
|
dtype=torch.bool
|
|
),
|
|
"rejected_mask": self.store.fetch_record(index, "rejected_mask").to(
|
|
dtype=torch.bool
|
|
),
|
|
}
|
|
|
|
|
|
@DatasetFactory.register("grpo")
|
|
class GRPODataset(BaseDataset):
|
|
"""Dataset for offline Group Relative Policy Optimization.
|
|
|
|
Each sample is one prompt with its group of responses and scalar
|
|
rewards — an independent training unit with no windowing or stride.
|
|
|
|
Expected storage layout (produced by JsonlStore or pre-tokenized):
|
|
|
|
- ``prompts``: List[Tensor] — one 1-D token tensor per record
|
|
- ``responses``: List[List[Tensor]] — G response tensors per record
|
|
- ``masks``: List[List[Tensor]] — G mask tensors per record
|
|
- ``rewards``: List[Tensor] — one 1-D float tensor (len G) per record
|
|
"""
|
|
|
|
required_keys = ["prompts", "responses", "masks", "rewards"]
|
|
|
|
def __getitem__(self, index: int) -> Dict[str, Tensor]:
|
|
prompts = self.store.fetch_record(index, "prompts")
|
|
responses = self.store.fetch_record(index, "responses")
|
|
masks = self.store.fetch_record(index, "masks")
|
|
rewards = self.store.fetch_record(index, "rewards")
|
|
return {
|
|
"prompts": prompts.to(dtype=torch.long),
|
|
"responses": [r.to(dtype=torch.long) for r in responses],
|
|
"masks": [m.to(dtype=torch.bool) for m in masks],
|
|
"rewards": rewards.to(dtype=torch.float32),
|
|
}
|