diff --git a/astrai/dataset/dataset.py b/astrai/dataset/dataset.py index 58051bb..03ebec7 100644 --- a/astrai/dataset/dataset.py +++ b/astrai/dataset/dataset.py @@ -346,7 +346,15 @@ class DatasetFactory(BaseFactory["BaseDataset"]): if processor is not None: store.load(load_path, processor=processor, **kwargs) else: - store.load(load_path, **kwargs) + load_kwargs = dict(kwargs) + if ( + tokenizer_path is not None + and storage_type == "jsonl" + and train_type in ("seq", "sft") + and "tokenizer_path" not in load_kwargs + ): + load_kwargs["tokenizer_path"] = tokenizer_path + store.load(load_path, **load_kwargs) return cls.create(train_type, store=store) diff --git a/astrai/dataset/storage.py b/astrai/dataset/storage.py index 81bf5af..fb4522e 100644 --- a/astrai/dataset/storage.py +++ b/astrai/dataset/storage.py @@ -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) diff --git a/tests/data/test_dataset.py b/tests/data/test_dataset.py index 29b12c2..bacab09 100644 --- a/tests/data/test_dataset.py +++ b/tests/data/test_dataset.py @@ -654,6 +654,92 @@ def test_jsonl_store_sft(base_test_env): assert item["loss_mask"].dtype == torch.bool +def test_sft_jsonl_default_messages_config(base_test_env): + """SFT loads a chat-style JSONL dir with no dataset_config.json. + + Falls back to the built-in messages config: every role except + ``assistant`` is masked, loss on assistant only. + """ + test_dir = base_test_env["test_dir"] + tokenizer = base_test_env["tokenizer"] + tokenizer.set_chat_template( + "{% for message in messages %}{{ message['role'] }}:{{ message['content'] }}\n{% endfor %}" + ) + tokenizer_path = _save_test_tokenizer(test_dir, tokenizer) + + data_dir = os.path.join(test_dir, "jsonl_data") + os.makedirs(data_dir, exist_ok=True) + records = [ + { + "messages": [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + ] + }, + { + "messages": [ + {"role": "user", "content": "bye"}, + {"role": "assistant", "content": "see you"}, + ] + }, + ] + with open(os.path.join(data_dir, "data.jsonl"), "w", encoding="utf-8") as f: + for record in records: + f.write(json.dumps(record, ensure_ascii=False) + "\n") + + dataset = DatasetFactory.load( + "sft", data_dir, window_size=8, tokenizer_path=tokenizer_path + ) + assert "sequence" in dataset.keys + assert "loss_mask" in dataset.keys + assert "position_ids" in dataset.keys + assert len(dataset) > 0 + item = dataset[0] + assert "input_ids" in item + assert "target_ids" in item + assert "loss_mask" in item + assert "position_ids" in item + assert item["loss_mask"].dtype == torch.bool + + +def test_sft_jsonl_explicit_config_takes_priority(base_test_env): + """When dataset_config.json exists, it overrides the default messages config.""" + test_dir = base_test_env["test_dir"] + tokenizer = base_test_env["tokenizer"] + tokenizer.set_chat_template( + "{% for message in messages %}{{ message['role'] }}:{{ message['content'] }}\n{% endfor %}" + ) + tokenizer_path = _save_test_tokenizer(test_dir, tokenizer) + + data_dir = _write_jsonl_dataset( + test_dir, + tokenizer_path, + [ + { + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + ] + } + ], + config_overrides={ + "input": { + "sections": [{"field": "messages", "action": "$role", "template": True}] + }, + "mask": {"user": "mask", "assistant": "train"}, + "mask_default": "mask", + "preprocessing": {"max_seq_len": 128}, + "output": {"position_ids_mode": "doc_reset"}, + }, + ) + dataset = DatasetFactory.load( + "sft", data_dir, window_size=8, tokenizer_path=tokenizer_path + ) + assert "sequence" in dataset.keys + assert "loss_mask" in dataset.keys + + def test_jsonl_store_pipeline_config_roundtrip(base_test_env): test_dir = base_test_env["test_dir"] config_path = os.path.join(test_dir, "dataset_config.json")