diff --git a/astrai/dataset/storage.py b/astrai/dataset/storage.py index be4a33e..e78575d 100644 --- a/astrai/dataset/storage.py +++ b/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) - if tokenizer_path is None: - raise ValueError( - f"JSONL dataset config must specify 'tokenizer_path': {config_path}" - ) - + tokenizer_path = raw_config.pop("tokenizer_path", None) or str(root) 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) - 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))) + for json_path in sorted(root.glob("*.json")): + if json_path.name == self.CONFIG_NAME: + continue + with open(json_path, "r", encoding="utf-8") as f: + try: + data = json.load(f) + except json.JSONDecodeError: + logger.warning("Failed to parse JSON file %s, skipping", json_path) + continue + if isinstance(data, list): + for item in data: + _process_item(item) + elif isinstance(data, dict): + _process_item(data) pos_ids = position_strategy.generate(doc_sequences) if pos_ids: diff --git a/astrai/preprocessing/pipeline.py b/astrai/preprocessing/pipeline.py index 5168802..a351fa9 100644 --- a/astrai/preprocessing/pipeline.py +++ b/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: - line = line.strip() - if not line: - continue - yield json.loads(line) + if path.endswith(".json"): + data = json.load(f) + if isinstance(data, dict): + yield data + 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(): diff --git a/tests/data/test_dataset.py b/tests/data/test_dataset.py index aaf0441..cd68e41 100644 --- a/tests/data/test_dataset.py +++ b/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"])