refactor: DatasetFactory.load accepts pre-built store instance

- load(store=...) binds directly, skipping format detection/processor
- load_path now optional when store is given
- Remove redundant from_store (merged into load)
- Caller can fully control Store construction + processor setup

astrai/dataset/dataset.py

@@ -365,28 +365,36 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
     def load(
         cls,
         train_type: str,
-        load_path: str,
+        load_path: Optional[str] = None,
-        window_size: int,
+        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* 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: Type of training dataset
-            load_path: Path to the data file
+            load_path: Path to the data file (ignored if *store* given)
             window_size: Window size for data sampling
             stride: Stride between consecutive samples (default: same as window_size)
             storage_type: Storage type ("h5", "bin", "jsonl") or None for auto-detection
-            tokenizer_path: Path to tokenizer. Used to build an on-the-fly
+            tokenizer_path: Path to tokenizer for lazy JSONL tokenisation
-                processor when loading raw JSONL with a record dataset
+            max_len: Max sequence length for the processor
-                (DPO/GRPO).  Ignored for pre-tokenised backends (H5/bin)
+            store: Pre-built, already-loaded Store instance
-                and for stream datasets (SEQ/SFT).
+            **kwargs: Extra arguments forwarded to ``dataset.load()``
-            max_len: Max sequence length for the processor.
-            **kwargs: Extra arguments forwarded to ``dataset.load()``.
 
         Returns:
             Loaded dataset instance
@@ -394,6 +402,14 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
         if stride is None:
             stride = window_size
 
+        if store is not None:
+            dataset = cls.create(train_type, window_size, stride)
+            dataset.storage = store
+            return dataset
+
+        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)
 
@@ -406,25 +422,6 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
 
         return dataset
 
-    @classmethod
-    def from_store(
-        cls,
-        train_type: str,
-        store: Store,
-        window_size: int = 0,
-        stride: Optional[int] = None,
-    ) -> "BaseDataset":
-        """Create a dataset bound to an already-loaded store.
-
-        The caller is responsible for constructing and loading the store
-        (including any processor).  The dataset simply wraps it.
-        """
-        if stride is None:
-            stride = window_size
-        dataset = cls.create(train_type, window_size, stride)
-        dataset.storage = store
-        return dataset
-
     @staticmethod
     def _maybe_build_processor(
         train_type: str,