From 8999ca89b84e2a0336b4357b7ef8bb4fe36b4cca Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sat, 4 Jul 2026 15:42:33 +0800 Subject: [PATCH] =?UTF-8?q?=EF=BB=BFfeat:=20add=20JSONL=20dataset=20store?= =?UTF-8?q?=20with=20on-the-fly=20tokenization?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add JsonlStore registered under "jsonl" in astrai/dataset/storage.py - Reuse PipelineConfig schema for JSONL dataset configuration - Update detect_format to recognize JSONL directories and files - Move save_h5/load_h5/save_bin/load_bin to astrai/serialization - Split astrai/serialization.py into checkpoint/dataset submodules - Add tests for JSONL detection, seq/SFT stores, and config roundtrip --- astrai/dataset/__init__.py | 4 + astrai/dataset/dataset.py | 16 +- astrai/dataset/storage.py | 178 +++++++++++------- astrai/preprocessing/writer.py | 2 +- astrai/serialization/__init__.py | 43 +++++ .../checkpoint.py} | 3 + astrai/serialization/dataset.py | 73 +++++++ tests/data/test_dataset.py | 140 ++++++++++++++ 8 files changed, 386 insertions(+), 73 deletions(-) create mode 100644 astrai/serialization/__init__.py rename astrai/{serialization.py => serialization/checkpoint.py} (99%) create mode 100644 astrai/serialization/dataset.py diff --git a/astrai/dataset/__init__.py b/astrai/dataset/__init__.py index cc8e7e4..562fa6f 100644 --- a/astrai/dataset/__init__.py +++ b/astrai/dataset/__init__.py @@ -5,10 +5,13 @@ from astrai.dataset.dataset import ( from astrai.dataset.sampler import ResumableDistributedSampler from astrai.dataset.storage import ( H5Store, + JsonlStore, MmapStore, Store, StoreFactory, detect_format, +) +from astrai.serialization import ( load_bin, load_h5, save_bin, @@ -22,6 +25,7 @@ __all__ = [ "StoreFactory", "H5Store", "MmapStore", + "JsonlStore", "detect_format", "save_h5", "load_h5", diff --git a/astrai/dataset/dataset.py b/astrai/dataset/dataset.py index 34e161d..4f07010 100644 --- a/astrai/dataset/dataset.py +++ b/astrai/dataset/dataset.py @@ -48,24 +48,26 @@ class BaseDataset(Dataset, ABC): f"Missing: {missing}" ) - def load(self, load_path: str, storage_type: Optional[str] = None): + def load(self, load_path: str, storage_type: Optional[str] = None, **kwargs): """Load dataset from the given path. Auto-detects the storage format if not specified. Args: load_path: Path to the data directory or file - storage_type: Force a specific storage type ("h5", "bin"), + storage_type: Force a specific storage type ("h5", "bin", "jsonl"), or None for auto-detection + **kwargs: Extra arguments forwarded to the store constructor and + to ``store.load()``. Raises: KeyError: If the loaded storage is missing required keys. """ if storage_type is None: storage_type = detect_format(load_path) - self.storage = StoreFactory.create(storage_type) + self.storage = StoreFactory.create(storage_type, **kwargs) self._load_path = load_path - self.storage.load(load_path) + self.storage.load(load_path, **kwargs) self._validate_keys() @property @@ -144,6 +146,7 @@ class DatasetFactory(BaseFactory["BaseDataset"]): window_size: int, stride: Optional[int] = None, storage_type: Optional[str] = None, + **kwargs, ) -> "BaseDataset": """Create and load a dataset in one step. @@ -152,7 +155,8 @@ class DatasetFactory(BaseFactory["BaseDataset"]): load_path: Path to the data file window_size: Window size for data sampling stride: Stride between consecutive samples (default: same as window_size) - storage_type: Storage type ("h5", "bin") or None for auto-detection + storage_type: Storage type ("h5", "bin", "jsonl") or None for auto-detection + **kwargs: Extra arguments forwarded to ``dataset.load()``. Returns: Loaded dataset instance @@ -161,7 +165,7 @@ class DatasetFactory(BaseFactory["BaseDataset"]): stride = window_size dataset = cls.create(train_type, window_size, stride) - dataset.load(load_path, storage_type=storage_type) + dataset.load(load_path, storage_type=storage_type, **kwargs) return dataset diff --git a/astrai/dataset/storage.py b/astrai/dataset/storage.py index adbad19..be4a33e 100644 --- a/astrai/dataset/storage.py +++ b/astrai/dataset/storage.py @@ -14,85 +14,31 @@ Key properties: - Explicit length: _length = min(total elements across keys), set at load, __len__ returns O(1) - Zero-copy mmap: MmapStore wraps np.memmap(mode="r"), all DataLoader - workers share OS page-cache pages + workers share OS page-cache pages """ import bisect import glob import json -import os +import logging from abc import ABC, abstractmethod from pathlib import Path from typing import Dict, List, Union -import h5py -import numpy as np import torch from torch import Tensor +from astrai.config.preprocess_config import PipelineConfig from astrai.factory import BaseFactory +from astrai.preprocessing.builder import MaskBuilderFactory +from astrai.preprocessing.position_id import PositionIdStrategyFactory +from astrai.serialization import ( + load_bin, + load_h5, +) +from astrai.tokenize import AutoTokenizer - -def save_h5(file_path: str, file_name: str, tensor_group: Dict[str, List[Tensor]]): - os.makedirs(file_path, exist_ok=True) - full_file_path = os.path.join(file_path, f"{file_name}.h5") - with h5py.File(full_file_path, "w") as f: - for key, tensors in tensor_group.items(): - grp = f.create_group(key) - for idx, tensor in enumerate(tensors): - arr = tensor.cpu().numpy() - grp.create_dataset(f"data_{idx}", data=arr) - - -def load_h5(file_path: str, share_memory=True) -> Dict[str, List[Tensor]]: - tensor_group: Dict[str, List[Tensor]] = {} - - root_path = Path(file_path) - h5_files = list(root_path.rglob("*.h5")) + list(root_path.rglob("*.hdf5")) - - for h5_file in h5_files: - with h5py.File(h5_file, "r") as f: - for key in f.keys(): - grp = f[key] - dsets = [] - for dset_name in grp.keys(): - dset = grp[dset_name] - tensor = torch.from_numpy(dset[:]) - if share_memory: - tensor = tensor.share_memory_() - dsets.append(tensor) - - if tensor_group.get(key) is None: - tensor_group[key] = [] - tensor_group[key].extend(dsets) - - return tensor_group - - -def save_bin(file_path: str, tensor_group: Dict[str, List[Tensor]]): - os.makedirs(file_path, exist_ok=True) - meta = {} - for key, tensors in tensor_group.items(): - cat = torch.cat(tensors, dim=0) - meta[key] = {"shape": list(cat.shape), "dtype": str(cat.dtype).split(".")[-1]} - np.asarray(cat.cpu().numpy()).tofile(os.path.join(file_path, f"{key}.bin")) - with open(os.path.join(file_path, "meta.json"), "w") as f: - json.dump(meta, f) - - -def load_bin(file_path: str) -> Dict[str, List[Tensor]]: - with open(os.path.join(file_path, "meta.json"), "r") as f: - meta = json.load(f) - segments: Dict[str, List[Tensor]] = {} - for key, info in meta.items(): - arr = np.memmap( - os.path.join(file_path, f"{key}.bin"), - dtype=info["dtype"], - mode="r+", - shape=tuple(info["shape"]), - ) - segments[key] = [torch.from_numpy(arr)] - return segments +logger = logging.getLogger(__name__) def detect_format(load_path: str) -> str: @@ -102,7 +48,7 @@ def detect_format(load_path: str) -> str: load_path: Directory or file path Returns: - Format string ("h5" or "bin") + Format string ("h5", "bin", or "jsonl") Raises: FileNotFoundError: If no supported data files are found @@ -112,6 +58,8 @@ def detect_format(load_path: str) -> str: suffix = root.suffix.lower() if suffix in (".h5", ".hdf5"): return "h5" + if suffix == ".jsonl": + return "jsonl" raise ValueError(f"Unsupported file format: {suffix}") h5_files = [ @@ -128,6 +76,11 @@ def detect_format(load_path: str) -> str: ) > 0 if has_meta: return "bin" + jsonl_files = [ + Path(p) for p in glob.glob(str(root / "**" / "*.jsonl"), recursive=True) + ] + if jsonl_files: + return "jsonl" raise FileNotFoundError(f"No supported data files found at {load_path}") @@ -264,3 +217,96 @@ class MmapStore(Store): self._normalize(all_raw) for tensors in self._data.values(): self._mmap_refs.extend(tensors) + + +@StoreFactory.register("jsonl") +class JsonlStore(Store): + """On-the-fly tokenization store for raw JSONL files. + + A JSONL dataset directory contains ``*.jsonl`` files plus a + ``dataset_config.json`` file that follows the same schema as + :class:`PipelineConfig` with an additional ``tokenizer_path`` field. + Records are tokenized when the store is loaded and concatenated into + segmented tensors matching the key layout expected by the dataset + classes (``sequence``, ``loss_mask``, ``position_ids``, ...). + """ + + CONFIG_NAME = "dataset_config.json" + + def load(self, path: str): + root = Path(path) + config_path = root / self.CONFIG_NAME + if not config_path.exists(): + raise FileNotFoundError( + f"JSONL dataset config not found: {config_path}. " + f"Expected {self.CONFIG_NAME} alongside *.jsonl files." + ) + + with open(config_path, "r", encoding="utf-8") as f: + raw_config = json.load(f) + + tokenizer_path = raw_config.pop("tokenizer_path", None) + if tokenizer_path is None: + raise ValueError( + f"JSONL dataset config must specify 'tokenizer_path': {config_path}" + ) + + self.config = PipelineConfig.from_dict(raw_config) + tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) + mask_builder = MaskBuilderFactory.create("sectioned") + position_strategy = PositionIdStrategyFactory.create( + self.config.output.position_ids_mode + ) + + raw: Dict[str, List[Tensor]] = {} + doc_sequences: List[List[int]] = [] + + for jsonl_path in sorted(root.glob("*.jsonl")): + with open(jsonl_path, "r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line: + continue + try: + item = json.loads(line) + except json.JSONDecodeError: + logger.warning( + "Failed to parse JSON line in %s, skipping", jsonl_path + ) + continue + + result = mask_builder.build(item, self.config, 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 key, ids in result.items(): + if key not in raw: + raw[key] = [] + raw[key].append(torch.tensor(ids, dtype=self._infer_dtype(ids))) + + pos_ids = position_strategy.generate(doc_sequences) + if pos_ids: + raw["position_ids"] = [torch.tensor(pos_ids, dtype=torch.int32)] + + self._normalize(raw) + + @staticmethod + def _primary_ids(result: dict) -> List[int]: + """Return the first integer list in *result* as the primary id sequence.""" + 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: + """Infer tensor dtype from the first element of a token/value list.""" + if ids and isinstance(ids[0], float): + return torch.float32 + return torch.int32 diff --git a/astrai/preprocessing/writer.py b/astrai/preprocessing/writer.py index ae74d5c..42488a6 100644 --- a/astrai/preprocessing/writer.py +++ b/astrai/preprocessing/writer.py @@ -14,8 +14,8 @@ from typing import Dict, List import torch -from astrai.dataset.storage import save_bin, save_h5 from astrai.factory import BaseFactory +from astrai.serialization import save_bin, save_h5 logger = logging.getLogger(__name__) diff --git a/astrai/serialization/__init__.py b/astrai/serialization/__init__.py new file mode 100644 index 0000000..e0e4a04 --- /dev/null +++ b/astrai/serialization/__init__.py @@ -0,0 +1,43 @@ +"""Serialization utilities for models and datasets. + +This package re-exports checkpoint helpers and dataset storage helpers so +that existing imports from ``astrai.serialization`` continue to work. +""" + +from astrai.serialization.checkpoint import ( + Checkpoint, + load_json, + load_model_config, + load_model_weights, + load_safetensors, + load_state_dict, + load_torch, + save_json, + save_model, + save_safetensors, + save_torch, +) +from astrai.serialization.dataset import ( + load_bin, + load_h5, + save_bin, + save_h5, +) + +__all__ = [ + "Checkpoint", + "load_json", + "load_model_config", + "load_model_weights", + "load_safetensors", + "load_state_dict", + "load_torch", + "save_json", + "save_model", + "save_safetensors", + "save_torch", + "load_bin", + "load_h5", + "save_bin", + "save_h5", +] diff --git a/astrai/serialization.py b/astrai/serialization/checkpoint.py similarity index 99% rename from astrai/serialization.py rename to astrai/serialization/checkpoint.py index 9be039c..0079f35 100644 --- a/astrai/serialization.py +++ b/astrai/serialization/checkpoint.py @@ -1,5 +1,8 @@ +"""Model checkpoint serialization helpers.""" + import io import json +import os import time from dataclasses import dataclass, field from pathlib import Path diff --git a/astrai/serialization/dataset.py b/astrai/serialization/dataset.py new file mode 100644 index 0000000..c106f40 --- /dev/null +++ b/astrai/serialization/dataset.py @@ -0,0 +1,73 @@ +"""Dataset storage serialization helpers (HDF5 / memory-mapped binary).""" + +import json +import os +from pathlib import Path +from typing import Dict, List + +import h5py +import numpy as np +import torch +from torch import Tensor + + +def save_h5(file_path: str, file_name: str, tensor_group: Dict[str, List[Tensor]]): + os.makedirs(file_path, exist_ok=True) + full_file_path = os.path.join(file_path, f"{file_name}.h5") + with h5py.File(full_file_path, "w") as f: + for key, tensors in tensor_group.items(): + grp = f.create_group(key) + for idx, tensor in enumerate(tensors): + arr = tensor.cpu().numpy() + grp.create_dataset(f"data_{idx}", data=arr) + + +def load_h5(file_path: str, share_memory=True) -> Dict[str, List[Tensor]]: + tensor_group: Dict[str, List[Tensor]] = {} + + root_path = Path(file_path) + h5_files = list(root_path.rglob("*.h5")) + list(root_path.rglob("*.hdf5")) + + for h5_file in h5_files: + with h5py.File(h5_file, "r") as f: + for key in f.keys(): + grp = f[key] + dsets = [] + for dset_name in grp.keys(): + dset = grp[dset_name] + tensor = torch.from_numpy(dset[:]) + if share_memory: + tensor = tensor.share_memory_() + dsets.append(tensor) + + if tensor_group.get(key) is None: + tensor_group[key] = [] + tensor_group[key].extend(dsets) + + return tensor_group + + +def save_bin(file_path: str, tensor_group: Dict[str, List[Tensor]]): + os.makedirs(file_path, exist_ok=True) + meta = {} + for key, tensors in tensor_group.items(): + cat = torch.cat(tensors, dim=0) + meta[key] = {"shape": list(cat.shape), "dtype": str(cat.dtype).split(".")[-1]} + np.asarray(cat.cpu().numpy()).tofile(os.path.join(file_path, f"{key}.bin")) + with open(os.path.join(file_path, "meta.json"), "w") as f: + json.dump(meta, f) + + +def load_bin(file_path: str) -> Dict[str, List[Tensor]]: + with open(os.path.join(file_path, "meta.json"), "r") as f: + meta = json.load(f) + segments: Dict[str, List[Tensor]] = {} + for key, info in meta.items(): + arr = np.memmap( + os.path.join(file_path, f"{key}.bin"), + dtype=info["dtype"], + mode="r+", + shape=tuple(info["shape"]), + ) + segments[key] = [torch.from_numpy(arr)] + return segments diff --git a/tests/data/test_dataset.py b/tests/data/test_dataset.py index 6f800dd..aaf0441 100644 --- a/tests/data/test_dataset.py +++ b/tests/data/test_dataset.py @@ -1,14 +1,18 @@ +import json import os import numpy as np import pytest import torch +from astrai.config.preprocess_config import PipelineConfig from astrai.dataset.dataset import DatasetFactory, SEQDataset from astrai.dataset.storage import ( H5Store, StoreFactory, detect_format, +) +from astrai.serialization import ( load_bin, save_bin, save_h5, @@ -19,6 +23,39 @@ def _rand_seq(length, vocab=1000): return torch.randint(0, vocab, (length,), dtype=torch.int64) +def _save_test_tokenizer(test_dir, tokenizer): + tokenizer_path = os.path.join(test_dir, "tokenizer") + os.makedirs(tokenizer_path, exist_ok=True) + tokenizer.save_pretrained(tokenizer_path) + return tokenizer_path + + +def _write_jsonl_dataset(test_dir, tokenizer_path, records, config_overrides=None): + data_dir = os.path.join(test_dir, "jsonl_data") + os.makedirs(data_dir, exist_ok=True) + + 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") + + config = { + "tokenizer_path": tokenizer_path, + "version": 1, + "input": {"sections": [{"field": "text", "action": "train"}]}, + "preprocessing": {"max_seq_len": 128}, + "output": {"position_ids_mode": "continuous"}, + } + if config_overrides: + config.update(config_overrides) + + with open( + os.path.join(data_dir, "dataset_config.json"), "w", encoding="utf-8" + ) as f: + json.dump(config, f, ensure_ascii=False, indent=2) + + return data_dir + + def _make_seq_dataset( test_dir, name="data", seq_length=200, train_type="seq", data=None, **load_kwargs ): @@ -372,3 +409,106 @@ def test_dataset_load_explicit_storage_type(base_test_env): dataset = _make_seq_dataset(test_dir, "explicit", storage_type="h5") assert len(dataset) > 0 assert dataset.count == 200 + + +def test_detect_format_jsonl_dir(base_test_env): + test_dir = base_test_env["test_dir"] + tokenizer_path = _save_test_tokenizer(test_dir, base_test_env["tokenizer"]) + data_dir = _write_jsonl_dataset( + test_dir, + tokenizer_path, + [{"text": "hello world"}, {"text": "foo bar baz"}], + ) + assert detect_format(data_dir) == "jsonl" + + +def test_jsonl_store_seq(base_test_env): + test_dir = base_test_env["test_dir"] + tokenizer_path = _save_test_tokenizer(test_dir, base_test_env["tokenizer"]) + data_dir = _write_jsonl_dataset( + test_dir, + tokenizer_path, + [{"text": "hello world"}, {"text": "foo bar baz qux"}], + config_overrides={"preprocessing": {"max_seq_len": 128, "min_chars": 0}}, + ) + + store = StoreFactory.create("jsonl") + store.load(data_dir) + assert len(store) > 0 + assert "sequence" in store.keys + + dataset = DatasetFactory.load("seq", data_dir, window_size=8) + assert len(dataset) > 0 + item = dataset[0] + assert "input_ids" in item + assert "target_ids" in item + assert item["input_ids"].dtype == torch.long + + +def test_jsonl_store_sft(base_test_env): + 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": "system", "content": "sys"}, + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + ] + } + ], + config_overrides={ + "input": { + "sections": [{"field": "messages", "action": "$role", "template": True}] + }, + "mask": {"system": "mask", "user": "mask", "assistant": "train"}, + "mask_default": "mask", + }, + ) + + store = StoreFactory.create("jsonl") + store.load(data_dir) + assert "sequence" in store.keys + assert "loss_mask" in store.keys + assert "position_ids" in store.keys + + dataset = DatasetFactory.load("sft", data_dir, window_size=8) + 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_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") + with open(config_path, "w", encoding="utf-8") as f: + json.dump( + { + "tokenizer_path": os.path.join(test_dir, "tokenizer"), + "version": 1, + "input": {"sections": [{"field": "text", "action": "train"}]}, + "mask": {"assistant": "train"}, + "preprocessing": {"max_seq_len": 64}, + "output": {"position_ids_mode": "doc_reset"}, + }, + f, + ensure_ascii=False, + indent=2, + ) + + with open(config_path, "r", encoding="utf-8") as f: + raw = json.load(f) + raw.pop("tokenizer_path") + config = PipelineConfig.from_dict(raw) + assert config.output.position_ids_mode == "doc_reset" + assert config.preprocessing.max_seq_len == 64