| """Load the JSONL files build_datasets.py produces into HF `datasets.Dataset` | |
| objects with the `images` column cast to lazily-loaded PIL images (paths on | |
| disk, only opened when a batch actually needs them — matters once the | |
| dataset is tens of thousands of rows).""" | |
| from __future__ import annotations | |
| from datasets import Dataset, Image, Sequence, load_dataset | |
| def load_jsonl_dataset(path: str, split: str | None = None) -> Dataset: | |
| ds = load_dataset("json", data_files=path)["train"] | |
| if split is not None: | |
| ds = ds.filter(lambda row: row["split"] == split) | |
| # "images" is a list of file-path strings in the JSONL; casting to | |
| # Sequence(Image()) makes datasets open them lazily as PIL images. | |
| ds = ds.cast_column("images", Sequence(Image())) | |
| return ds | |
| def dataset_stats(ds: Dataset) -> dict: | |
| rounds = sorted(set(ds["round"])) if "round" in ds.column_names else [] | |
| return {"n_rows": len(ds), "n_rounds": len(rounds), "rounds": rounds} | |