feat: support SFT directly from JSONL without dataset_config.json
- JsonlStore falls back to built-in messages config when no config file found and tokenizer_path is provided - DatasetFactory.load forwards tokenizer_path to store for SFT/SEQ+jsonl - assistant turns train, other roles masked, position_ids doc_reset
This commit is contained in:
+42
-14
@@ -56,6 +56,7 @@ from typing import Callable, Dict, List, Optional, Tuple, Union
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
from astrai.config.preprocess_config import PipelineConfig
|
||||
from astrai.factory import BaseFactory
|
||||
from astrai.preprocessing.transform import TokenizeTransform
|
||||
from astrai.serialization import (
|
||||
@@ -545,18 +546,29 @@ class JsonlSource:
|
||||
|
||||
@StoreFactory.register("jsonl")
|
||||
class JsonlStore(Store, Streamable, Recordable):
|
||||
"""JSONL reader with two tokenisation modes.
|
||||
"""JSONL reader with eager/lazy tokenisation modes.
|
||||
|
||||
A JSONL dataset is a ``.jsonl`` file or a directory of ``*.jsonl``
|
||||
files plus (optionally) a ``dataset_config.json`` describing the
|
||||
tokenization pipeline.
|
||||
|
||||
Two modes, selected at :meth:`load` time:
|
||||
Three ways to supply an eager transform (first match wins):
|
||||
|
||||
- **Eager** (default): applies a :class:`TokenizeTransform` to every
|
||||
record at load time and registers per-key tensors via
|
||||
``_normalize``. Both ``fetch`` (stream) and ``fetch_record``
|
||||
(record) work.
|
||||
- **Explicit** (``transform=``): caller-built
|
||||
:class:`TokenizeTransform` applied eagerly.
|
||||
- **Config file**: ``dataset_config.json`` alongside the ``*.jsonl``
|
||||
files — loaded via :meth:`TokenizeTransform.from_config_file`.
|
||||
- **Default messages** (``tokenizer_path=`` given, no config file):
|
||||
a built-in chatml config that tokenises the ``messages`` field,
|
||||
masking every role except ``assistant`` (loss on assistant only).
|
||||
Lets SFT/SEQ train straight from a chat-style JSONL directory
|
||||
without a hand-written config.
|
||||
|
||||
Two tokenisation modes, selected at :meth:`load` time:
|
||||
|
||||
- **Eager** (default): applies the transform to every record at load
|
||||
time and registers per-key tensors via ``_normalize``. Both
|
||||
``fetch`` (stream) and ``fetch_record`` (record) work.
|
||||
- **Lazy** (``processor=fn`` passed): keeps raw records and defers
|
||||
tokenisation to ``fetch_record``. Only record access works —
|
||||
``len(store)`` returns ``num_records``; stream primitives raise.
|
||||
@@ -565,6 +577,16 @@ class JsonlStore(Store, Streamable, Recordable):
|
||||
CONFIG_NAME = "dataset_config.json"
|
||||
segments_are_records = True
|
||||
|
||||
_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 __init__(
|
||||
self,
|
||||
window_size: int = 0,
|
||||
@@ -587,14 +609,20 @@ class JsonlStore(Store, Streamable, Recordable):
|
||||
if transform is None:
|
||||
root = Path(path)
|
||||
config_path = root / self.CONFIG_NAME if root.is_dir() else None
|
||||
if config_path is None or not config_path.exists():
|
||||
raise FileNotFoundError(
|
||||
f"JSONL dataset config not found. Expected "
|
||||
f"{self.CONFIG_NAME} alongside *.jsonl files, pass an "
|
||||
f"explicit transform, or pass processor= for lazy "
|
||||
f"on-the-fly tokenisation."
|
||||
)
|
||||
transform = TokenizeTransform.from_config_file(str(config_path))
|
||||
if config_path is not None and config_path.exists():
|
||||
transform = TokenizeTransform.from_config_file(str(config_path))
|
||||
else:
|
||||
tokenizer_path = kwargs.get("tokenizer_path")
|
||||
if not tokenizer_path:
|
||||
raise FileNotFoundError(
|
||||
f"JSONL dataset config not found. Expected "
|
||||
f"{self.CONFIG_NAME} alongside *.jsonl files, pass an "
|
||||
f"explicit transform, pass processor= for lazy "
|
||||
f"on-the-fly tokenisation, or pass tokenizer_path= to "
|
||||
f"use the built-in messages config."
|
||||
)
|
||||
config = PipelineConfig.from_dict(self._DEFAULT_MESSAGES_CONFIG)
|
||||
transform = TokenizeTransform(config, tokenizer_path)
|
||||
|
||||
transformed = transform.apply(records)
|
||||
self._normalize(transformed)
|
||||
|
||||
Reference in New Issue
Block a user