refactor: break JsonlStore→preprocessing circular dependency

- 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
This commit is contained in:
2026-08-07 23:22:32 +08:00
parent 1b1f1a0707
commit f163520fff
3 changed files with 51 additions and 43 deletions
+4 -29
View File
@@ -55,9 +55,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 (
load_bin,
load_bin_offsets,
@@ -536,19 +534,8 @@ class JsonlStore(Store, Streamable, Recordable):
``len(store)`` returns ``num_records``; stream primitives raise.
"""
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,
@@ -569,22 +556,10 @@ class JsonlStore(Store, Streamable, Recordable):
return
if transform is None:
root = Path(path)
config_path = root / self.CONFIG_NAME if root.is_dir() else None
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)
raise ValueError(
"JsonlStore eager mode requires transform=. "
"Use DatasetFactory.load() which auto-constructs it."
)
transformed = transform.apply(records)
self._normalize(transformed)