diff --git a/deepspec/data/jsonl_dataset.py b/deepspec/data/jsonl_dataset.py index 65e11ec4..05f0219a 100644 --- a/deepspec/data/jsonl_dataset.py +++ b/deepspec/data/jsonl_dataset.py @@ -101,21 +101,22 @@ def _build_all_line_starts(self): self.num_data_per_file.append(len(starts)) continue - handle = open(path, "rb") - mm = mmap.mmap(handle.fileno(), 0, access=mmap.ACCESS_READ) - self.files[idx] = handle - self.mmaps[idx] = mm starts = [] - mm.seek(0) - pos = 0 - while True: - starts.append(pos) - line = mm.readline() - if not line: - break - pos = mm.tell() - if starts and mm.size() == pos: - starts.pop() + if os.path.getsize(path) > 0: + handle = open(path, "rb") + mm = mmap.mmap(handle.fileno(), 0, access=mmap.ACCESS_READ) + self.files[idx] = handle + self.mmaps[idx] = mm + mm.seek(0) + pos = 0 + while True: + starts.append(pos) + line = mm.readline() + if not line: + break + pos = mm.tell() + if starts and mm.size() == pos: + starts.pop() self.line_starts_per_file[idx] = starts self.num_data_per_file.append(len(starts)) if cache_path is not None: diff --git a/tests/test_jsonl_dataset.py b/tests/test_jsonl_dataset.py new file mode 100644 index 00000000..3a51e935 --- /dev/null +++ b/tests/test_jsonl_dataset.py @@ -0,0 +1,65 @@ +import importlib.util +import json +import sys +import types +from pathlib import Path + + +def _load_jsonl_dataset_module(monkeypatch, cache_dir: Path): + class Dataset: + pass + + torch = types.ModuleType("torch") + torch.utils = types.SimpleNamespace(data=types.SimpleNamespace(Dataset=Dataset)) + constants = types.ModuleType("deepspec.utils.constant") + constants.CACHE_DIR = str(cache_dir) + tqdm_module = types.ModuleType("tqdm") + tqdm_module.tqdm = lambda iterable, **_kwargs: iterable + + monkeypatch.setitem(sys.modules, "torch", torch) + monkeypatch.setitem(sys.modules, "deepspec", types.ModuleType("deepspec")) + monkeypatch.setitem(sys.modules, "deepspec.utils", types.ModuleType("deepspec.utils")) + monkeypatch.setitem(sys.modules, "deepspec.utils.constant", constants) + monkeypatch.setitem(sys.modules, "tqdm", tqdm_module) + + module_path = Path(__file__).parents[1] / "deepspec" / "data" / "jsonl_dataset.py" + spec = importlib.util.spec_from_file_location("jsonl_dataset_under_test", module_path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_jsonl_dataset_skips_empty_shards(monkeypatch, tmp_path): + module = _load_jsonl_dataset_module(monkeypatch, tmp_path / "cache") + empty_path = tmp_path / "empty.jsonl" + data_path = tmp_path / "data.jsonl" + empty_path.write_bytes(b"") + rows = [{"id": 1}, {"id": 2}] + data_path.write_text("".join(json.dumps(row) + "\n" for row in rows), encoding="utf-8") + + dataset = module.JsonLineDataset([data_path, empty_path]) + + assert len(dataset) == 2 + assert [dataset[index] for index in range(len(dataset))] == rows + assert dataset.num_data_per_file == [2, 0] + dataset.close() + + cached_dataset = module.JsonLineDataset([data_path, empty_path]) + + assert len(cached_dataset) == 2 + assert cached_dataset.num_data_per_file == [2, 0] + cached_dataset.close() + + +def test_jsonl_dataset_accepts_only_empty_shards(monkeypatch, tmp_path): + module = _load_jsonl_dataset_module(monkeypatch, tmp_path / "cache") + paths = [tmp_path / "empty-1.jsonl", tmp_path / "empty-2.jsonl"] + for path in paths: + path.write_bytes(b"") + + dataset = module.JsonLineDataset(paths) + + assert len(dataset) == 0 + assert dataset.num_data_per_file == [0, 0] + dataset.close()