feat: support raw JSON files in dataset pipeline and JsonlStore

- detect_format now recognizes .json directories as jsonl store
- JsonlStore loads .json arrays and dicts alongside .jsonl
- tokenizer_path defaults to dataset dir when omitted
- Pipeline._iter_items handles .json files (arrays/single dict)
- Tests: detect_format, seq load, self-contained dataset dir

astrai/dataset/storage.py

@@ -81,6 +81,11 @@ def detect_format(load_path: str) -> str:
     ]
     if jsonl_files:
         return "jsonl"
+    json_files = [
+        Path(p) for p in glob.glob(str(root / "**" / "*.json"), recursive=True)
+    ]
+    if json_files:
+        return "jsonl"
     raise FileNotFoundError(f"No supported data files found at {load_path}")
 
 
@@ -245,12 +250,7 @@ class JsonlStore(Store):
         with open(config_path, "r", encoding="utf-8") as f:
             raw_config = json.load(f)
 
-        tokenizer_path = raw_config.pop("tokenizer_path", None)
+        tokenizer_path = raw_config.pop("tokenizer_path", None) or str(root)
-        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")
@@ -261,6 +261,21 @@ class JsonlStore(Store):
         raw: Dict[str, List[Tensor]] = {}
         doc_sequences: List[List[int]] = []
 
+        def _process_item(item: dict) -> None:
+            nonlocal raw, doc_sequences
+            result = mask_builder.build(item, self.config, tokenizer)
+            if result is None:
+                return
+            result.pop("domain", None)
+            primary_ids = self._primary_ids(result)
+            if not primary_ids:
+                return
+            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)))
+
         for jsonl_path in sorted(root.glob("*.jsonl")):
             with open(jsonl_path, "r", encoding="utf-8") as f:
                 for line in f:
@@ -274,21 +289,22 @@ class JsonlStore(Store):
                             "Failed to parse JSON line in %s, skipping", jsonl_path
                         )
                         continue
+                    _process_item(item)
 
-                    result = mask_builder.build(item, self.config, tokenizer)
+        for json_path in sorted(root.glob("*.json")):
-                    if result is None:
+            if json_path.name == self.CONFIG_NAME:
-                        continue
+                continue
-
+            with open(json_path, "r", encoding="utf-8") as f:
-                    result.pop("domain", None)
+                try:
-                    primary_ids = self._primary_ids(result)
+                    data = json.load(f)
-                    if not primary_ids:
+                except json.JSONDecodeError:
-                        continue
+                    logger.warning("Failed to parse JSON file %s, skipping", json_path)
-
+                    continue
-                    doc_sequences.append(primary_ids)
+            if isinstance(data, list):
-                    for key, ids in result.items():
+                for item in data:
-                        if key not in raw:
+                    _process_item(item)
-                            raw[key] = []
+            elif isinstance(data, dict):
-                        raw[key].append(torch.tensor(ids, dtype=self._infer_dtype(ids)))
+                _process_item(data)
 
         pos_ids = position_strategy.generate(doc_sequences)
         if pos_ids:

astrai/preprocessing/pipeline.py

@@ -149,11 +149,18 @@ class Pipeline:
     def _iter_items(self):
         for path in self.paths:
             with open(path, "r", encoding="utf-8") as f:
-                for line in f:
+                if path.endswith(".json"):
-                    line = line.strip()
+                    data = json.load(f)
-                    if not line:
+                    if isinstance(data, dict):
-                        continue
+                        yield data
-                    yield json.loads(line)
+                    elif isinstance(data, list):
+                        yield from data
+                else:
+                    for line in f:
+                        line = line.strip()
+                        if not line:
+                            continue
+                        yield json.loads(line)
 
     def _flush(self, domains, shard_idx):
         for domain, keys in domains.items():

tests/data/test_dataset.py

@@ -411,6 +411,32 @@ def test_dataset_load_explicit_storage_type(base_test_env):
     assert dataset.count == 200
 
 
+def _write_json_dataset(test_dir, tokenizer_path, records, config_overrides=None):
+    """Write JSON (not JSONL) dataset — array of objects."""
+    data_dir = os.path.join(test_dir, "json_data")
+    os.makedirs(data_dir, exist_ok=True)
+
+    with open(os.path.join(data_dir, "data.json"), "w", encoding="utf-8") as f:
+        json.dump(records, f, ensure_ascii=False)
+
+    config = {
+        "tokenizer_path": tokenizer_path,
+        "version": 1,
+        "input": {"sections": [{"field": "text", "action": "train"}]},
+        "preprocessing": {"max_seq_len": 128, "min_chars": 0},
+        "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 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"])
@@ -422,6 +448,89 @@ def test_detect_format_jsonl_dir(base_test_env):
     assert detect_format(data_dir) == "jsonl"
 
 
+def test_detect_format_json_dir(base_test_env):
+    """detect_format returns 'jsonl' for directory with .json files."""
+    test_dir = base_test_env["test_dir"]
+    tokenizer_path = _save_test_tokenizer(test_dir, base_test_env["tokenizer"])
+    data_dir = _write_json_dataset(
+        test_dir,
+        tokenizer_path,
+        [{"text": "hello world"}, {"text": "foo bar baz qux"}],
+    )
+    assert detect_format(data_dir) == "jsonl"
+
+
+def test_json_store_seq(base_test_env):
+    """JsonlStore loads .json array correctly."""
+    test_dir = base_test_env["test_dir"]
+    tokenizer_path = _save_test_tokenizer(test_dir, base_test_env["tokenizer"])
+    data_dir = _write_json_dataset(
+        test_dir,
+        tokenizer_path,
+        [{"text": "hello world"}, {"text": "foo bar baz qux"}],
+    )
+
+    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
+
+
+def test_json_store_no_tokenizer_path(base_test_env):
+    """JsonlStore uses dataset dir as tokenizer_path when omitted."""
+    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 %}"
+    )
+
+    data_dir = os.path.join(test_dir, "self_contained")
+    os.makedirs(data_dir, exist_ok=True)
+
+    # Save tokenizer files directly in the dataset directory
+    tokenizer.save_pretrained(data_dir)
+
+    # Write .json data
+    records = [
+        {
+            "messages": [
+                {"role": "user", "content": "hi"},
+                {"role": "assistant", "content": "hello"},
+            ]
+        }
+    ]
+    with open(os.path.join(data_dir, "data.json"), "w", encoding="utf-8") as f:
+        json.dump(records, f, ensure_ascii=False)
+
+    # dataset_config.json WITHOUT tokenizer_path
+    config = {
+        "version": 1,
+        "input": {
+            "sections": [{"field": "messages", "action": "$role", "template": True}]
+        },
+        "mask": {"user": "mask", "assistant": "train"},
+        "mask_default": "mask",
+        "preprocessing": {"max_seq_len": 128, "min_chars": 0},
+        "output": {"position_ids_mode": "continuous"},
+    }
+    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)
+
+    store = StoreFactory.create("jsonl")
+    store.load(data_dir)
+    assert len(store) > 0
+    assert "sequence" in store.keys
+    assert "loss_mask" in store.keys
+
+
 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"])