From 613a0460207c18ddd4e2cff3381ad7291bcdf8e7 Mon Sep 17 00:00:00 2001 From: Urmil Chandarana Date: Wed, 16 Sep 2026 15:02:36 -0700 Subject: [PATCH 1/3] FEATURE: Added supporting loader for Image-Text-To-Text Datasets on HuggingFace --- README.md | 40 ++++++ agml/__init__.py | 4 +- agml/data/__init__.py | 1 + agml/data/image_text_to_text_loader.py | 191 +++++++++++++++++++++++++ 4 files changed, 235 insertions(+), 1 deletion(-) create mode 100644 agml/data/image_text_to_text_loader.py diff --git a/README.md b/README.md index 91fb4a70..06928b41 100644 --- a/README.md +++ b/README.md @@ -80,6 +80,46 @@ trainer = Trainer(model=model, train_dataset=dataset["train"], eval_dataset=data trainer.train() ``` +### Image-Text-to-Text Datasets + +Some AgML datasets (e.g. `Project-AgML/AgroOmni`) are too large to fit the standard `load_dataset()` +archive handling, which fully extracts every archive to disk before reading it. These datasets instead store images +as raw relative paths inside zip shards, alongside a single shared `path -> shard` index, so no extraction step is +needed and no per-archive metadata has to be duplicated. Because of that, they aren't `load_dataset()`-compatible, +and are loaded with `loadImageTextToTextDataset` instead: + +```python +from agml import loadImageTextToTextDataset + +# Whole dataset, split-aware: +ds = loadImageTextToTextDataset( + "Project-AgML/AgroOmni", + token=HF_TOKEN +) + +ds["train"][0] # {'id', 'messages', 'raw_metadata', 'images': [PIL.Image, ...]} + +# Single split: +train_ds = loadImageTextToTextDataset( + "Project-AgML/AgroOmni", + split="train", + token=HF_TOKEN +) +``` + +`loadImageTextToTextDataset(repo_id, split=None, cache_dir=None, token=None)` returns a `DatasetDict` when `split` +is omitted, or a single `Dataset` when a split name is given. Splits are auto-discovered from the metadata parquet +filenames (e.g. `train-0000-of-0001.parquet` becomes `train`), the same convention `load_dataset()` itself uses; if +no split can be inferred, every row is placed under a single `train` split. + +Image bytes are only read the first time a row is actually accessed (`ds["train"][0]`, a slice, or a batch), never +eagerly over the whole dataset, and repeated access to the same image is cached. The `ImageTextToTextShardStore` +backing that lazy decoding is attached to the returned dataset as `ds.store`, in case you want to inspect its cache, +call `ds.store.getBytes(path)` directly, or free its open shard file handles (`del ds.store`). + +If you're using this with a PyTorch `DataLoader`, prefer `num_workers > 0` so each worker process opens its own +shard handles rather than sharing one across threads. + ## Public Datasets AgML contains a wide variety of public datasets from various locations across the world: diff --git a/agml/__init__.py b/agml/__init__.py index d795fe5b..a8fff9f9 100644 --- a/agml/__init__.py +++ b/agml/__init__.py @@ -38,5 +38,7 @@ def _setup(): del _setup # noqa -# There are no top-level imported functions or classes, only the modules. +# Besides the modules, `loadImageTextToTextDataset` is exposed at the top level +# since it is the entry point for loading AgML image-text-to-text datasets from the Hub. from . import backend, data, io, synthetic, viz +from .data import loadImageTextToTextDataset diff --git a/agml/data/__init__.py b/agml/data/__init__.py index 2a7ac751..9c42db55 100644 --- a/agml/data/__init__.py +++ b/agml/data/__init__.py @@ -20,3 +20,4 @@ from .tools import coco_to_bboxes, convert_bbox_format from .hf_loader import HuggingFaceDataLoader from .multispectral_hf_loader import MultispectralDataLoader +from .image_text_to_text_loader import ImageTextToTextShardStore, loadImageTextToTextDataset diff --git a/agml/data/image_text_to_text_loader.py b/agml/data/image_text_to_text_loader.py new file mode 100644 index 00000000..be74fc5a --- /dev/null +++ b/agml/data/image_text_to_text_loader.py @@ -0,0 +1,191 @@ +import io +import os +import re +import zipfile +from collections import defaultdict +from functools import lru_cache +from typing import Optional, Union + +import pandas as pd +from PIL import Image + +try: + from datasets import Dataset, DatasetDict +except ImportError: + raise ImportError( + "The `datasets` library is required to use loadImageTextToTextDataset. " + "Please install it using `pip install datasets`." + ) + +try: + from huggingface_hub import snapshot_download +except ImportError: + raise ImportError( + "The `huggingface_hub` library is required to use loadImageTextToTextDataset. " + "Please install it using `pip install huggingface_hub`." + ) + +# Matches metadata parquet filenames such as "train-0000-of-0001.parquet", +# capturing the split name so splits can be auto-discovered. +_SPLIT_PARQUET_PATTERN = re.compile(r"^(?P.+)-\d+-of-\d+\.parquet$") + +# The relative path (from the repo root) where image shards and the shard index live. +_IMAGES_DIRNAME = "images" + +# The parquet file mapping each image path to the shard file that contains it. +_SHARD_INDEX_FILENAME = "path_to_shard.parquet" + +# The split name used when no split can be inferred from the metadata filenames. +_DEFAULT_SPLIT_NAME = "train" + + +class ImageTextToTextShardStore: + """ + Looks up and reads image bytes out of image-text-to-text shards without extracting them to disk. + + On construction, this reads `path_to_shard.parquet` into an in-memory + `{path: shard_file}` index. Each shard's `zipfile.ZipFile` handle is opened + lazily on first access and kept open for reuse, and reads are cached with + an `lru_cache` since the same image path is often referenced by multiple rows. + """ + + def __init__(self, images_dir: str): + self.images_dir = images_dir + index_path = os.path.join(images_dir, _SHARD_INDEX_FILENAME) + index_df = pd.read_parquet(index_path) + self._path_to_shard = dict(zip(index_df["path"], index_df["shard_file"])) + self._zip_cache = {} + + def _getZip(self, shard_file: str): + """Returns the open `ZipFile` handle for a shard, opening it on first access.""" + zf = self._zip_cache.get(shard_file) + if zf is None: + shard_path = os.path.join(self.images_dir, shard_file) + zf = zipfile.ZipFile(shard_path, "r") + self._zip_cache[shard_file] = zf + return zf + + @lru_cache(maxsize=4096) + def _getBytesCached(self, path: str) -> bytes: + """Reads raw image bytes for a path, caching repeat reads of the same path.""" + shard_file = self._path_to_shard[path] + return self._getZip(shard_file).read(path) + + def getBytes(self, path: str) -> bytes: + """Public entry point for reading the raw bytes of a single image path.""" + return self._getBytesCached(path) + + def __del__(self): + # Best effort cleanup, the process may already be tearing down. + for zf in getattr(self, "_zip_cache", {}).values(): + try: + zf.close() + except Exception: + pass + + +def _makeDecodeTransform(store: ImageTextToTextShardStore): + """Builds the lazy `set_transform` callback that decodes the `images` column via `store`.""" + + def decode(batch): + decoded_rows = [] + for images in batch["images"]: + decoded_row = [] + for image_ref in images: + path = image_ref["path"] if isinstance(image_ref, dict) else image_ref + raw_bytes = store.getBytes(path) + decoded_row.append(Image.open(io.BytesIO(raw_bytes))) + decoded_rows.append(decoded_row) + batch["images"] = decoded_rows + return batch + + return decode + + +def _discoverSplits(repo_dir: str) -> dict: + """ + Groups metadata parquet files in `repo_dir` by their auto-discovered split name. + + Filenames matching `-N-of-M.parquet` are grouped under ``. If + none of the metadata parquet files match that convention, every parquet file + at the repo root is treated as a single `train` split instead of failing. + """ + split_to_files = defaultdict(list) + all_parquet_files = [] + for filename in sorted(os.listdir(repo_dir)): + if not filename.endswith(".parquet"): + continue + all_parquet_files.append(os.path.join(repo_dir, filename)) + + match = _SPLIT_PARQUET_PATTERN.match(filename) + if match is not None: + split_to_files[match.group("split")].append(os.path.join(repo_dir, filename)) + + if not split_to_files and all_parquet_files: + split_to_files[_DEFAULT_SPLIT_NAME] = all_parquet_files + + if not split_to_files: + raise FileNotFoundError(f"No metadata parquet files were found in {repo_dir}.") + + return split_to_files + + +def loadImageTextToTextDataset( + repo_id: str, + split: Optional[str] = None, + cache_dir: Optional[str] = None, + token: Optional[str] = None, +) -> Union["DatasetDict", "Dataset"]: + """ + Loads an AgML image-text-to-text shard dataset from the Hub. + + This downloads `repo_id` with `snapshot_download`, builds an + `ImageTextToTextShardStore` over the downloaded `images/` shards, and loads + the metadata parquet(s) into an Arrow-backed dataset with a lazy transform + attached that decodes the `images` column into `PIL.Image` objects on + access. Splits are auto-discovered from the metadata filenames (e.g. + `train-0000-of-0001.parquet` becomes `train`); if no split can be inferred, + every row is placed under a single `train` split. + + The backing `ImageTextToTextShardStore` is attached to the returned dataset + as a `.store` attribute rather than returned separately, so callers can + inspect its cache or free its open file handles (`del ds.store`) without + changing the common `ds = loadImageTextToTextDataset(...)` call shape. + + Parameters + ---------- + repo_id : str + The Hugging Face Hub dataset repo id, e.g. `Project-AgML/AgroOmni-plainziptest`. + split : str, optional + If given, only that split is returned (as a `Dataset` rather than a `DatasetDict`). + cache_dir : str, optional + The local directory to cache the downloaded repo in. + token : str, optional + A Hugging Face Hub token, needed for private or gated repos. + + Returns + ------- + DatasetDict or Dataset + A `DatasetDict` (or a single `Dataset` if `split` was given), with the + `ImageTextToTextShardStore` backing the lazy image decoding attached as + its `.store` attribute. + """ + repo_dir = snapshot_download(repo_id=repo_id, repo_type="dataset", cache_dir=cache_dir, token=token) + images_dir = os.path.join(repo_dir, _IMAGES_DIRNAME) + store = ImageTextToTextShardStore(images_dir) + decode = _makeDecodeTransform(store) + + split_to_files = _discoverSplits(repo_dir) + dataset_dict = DatasetDict() + for split_name, parquet_files in split_to_files.items(): + split_ds = Dataset.from_parquet(parquet_files) + split_ds.set_transform(decode) + dataset_dict[split_name] = split_ds + + if split is not None: + result = dataset_dict[split] + result.store = store + return result + + dataset_dict.store = store + return dataset_dict From 31a163c3d163117157eb34dfba995946a6d98932 Mon Sep 17 00:00:00 2001 From: Urmil Chandarana Date: Wed, 16 Sep 2026 16:17:12 -0700 Subject: [PATCH 2/3] BUGFIX: Added fixes for dataset configs and clean implementation for downloading requested configs only --- README.md | 17 ++- agml/data/image_text_to_text_loader.py | 146 +++++++++++++++++++++---- 2 files changed, 140 insertions(+), 23 deletions(-) diff --git a/README.md b/README.md index 06928b41..0c7fa1e6 100644 --- a/README.md +++ b/README.md @@ -105,12 +105,21 @@ train_ds = loadImageTextToTextDataset( split="train", token=HF_TOKEN ) + +# Multi-config repo (each config under its own folder): pass `config` to select one. +ds = loadImageTextToTextDataset( + "Project-AgML/MIRAGE", + config="MMST_Standard", + token=HF_TOKEN, +) ``` -`loadImageTextToTextDataset(repo_id, split=None, cache_dir=None, token=None)` returns a `DatasetDict` when `split` -is omitted, or a single `Dataset` when a split name is given. Splits are auto-discovered from the metadata parquet -filenames (e.g. `train-0000-of-0001.parquet` becomes `train`), the same convention `load_dataset()` itself uses; if -no split can be inferred, every row is placed under a single `train` split. +`loadImageTextToTextDataset(repo_id, config=None, split=None, cache_dir=None, token=None)` returns a `DatasetDict` +when `split` is omitted, or a single `Dataset` when a split name is given. Splits are auto-discovered from the +metadata parquet filenames (e.g. `train-0000-of-0001.parquet` becomes `train`), the same convention `load_dataset()` +itself uses; if no split can be inferred, every row is placed under a single `train` split. For repos with multiple +configs, each living under its own folder, `config` is required, if it's omitted the error message lists the +configs available in that repo. Image bytes are only read the first time a row is actually accessed (`ds["train"][0]`, a slice, or a batch), never eagerly over the whole dataset, and repeated access to the same image is cached. The `ImageTextToTextShardStore` diff --git a/agml/data/image_text_to_text_loader.py b/agml/data/image_text_to_text_loader.py index be74fc5a..7ae3e3cb 100644 --- a/agml/data/image_text_to_text_loader.py +++ b/agml/data/image_text_to_text_loader.py @@ -7,6 +7,7 @@ from typing import Optional, Union import pandas as pd +import yaml from PIL import Image try: @@ -18,7 +19,8 @@ ) try: - from huggingface_hub import snapshot_download + from huggingface_hub import hf_hub_download, list_repo_files, snapshot_download + from huggingface_hub.utils import EntryNotFoundError except ImportError: raise ImportError( "The `huggingface_hub` library is required to use loadImageTextToTextDataset. " @@ -38,6 +40,9 @@ # The split name used when no split can be inferred from the metadata filenames. _DEFAULT_SPLIT_NAME = "train" +# The dataset card holding the `configs` YAML front matter that names the default config. +_DATASET_CARD_FILENAME = "README.md" + class ImageTextToTextShardStore: """ @@ -102,36 +107,117 @@ def decode(batch): return decode -def _discoverSplits(repo_dir: str) -> dict: +def _discoverSplits(config_dir: str) -> dict: """ - Groups metadata parquet files in `repo_dir` by their auto-discovered split name. + Groups metadata parquet files in `config_dir` by their auto-discovered split name. Filenames matching `-N-of-M.parquet` are grouped under ``. If none of the metadata parquet files match that convention, every parquet file - at the repo root is treated as a single `train` split instead of failing. + in `config_dir` is treated as a single `train` split instead of failing. """ split_to_files = defaultdict(list) all_parquet_files = [] - for filename in sorted(os.listdir(repo_dir)): + for filename in sorted(os.listdir(config_dir)): if not filename.endswith(".parquet"): continue - all_parquet_files.append(os.path.join(repo_dir, filename)) + all_parquet_files.append(os.path.join(config_dir, filename)) match = _SPLIT_PARQUET_PATTERN.match(filename) if match is not None: - split_to_files[match.group("split")].append(os.path.join(repo_dir, filename)) + split_to_files[match.group("split")].append(os.path.join(config_dir, filename)) if not split_to_files and all_parquet_files: split_to_files[_DEFAULT_SPLIT_NAME] = all_parquet_files if not split_to_files: - raise FileNotFoundError(f"No metadata parquet files were found in {repo_dir}.") + raise FileNotFoundError(f"No metadata parquet files were found in {config_dir}.") return split_to_files +def _fetchDefaultConfigName(repo_id: str, token: Optional[str]) -> Optional[str]: + """ + Reads the default config name out of the dataset card's `configs` YAML front matter. + + This mirrors how `datasets`/the Hub itself pick a default config: a repo's + `README.md` starts with a `---`-delimited YAML block listing each config + under `configs`, with the one meant to load by default marked `default: true`. + If no entry is marked, the first one listed is used, matching Hub behavior. + + Only `README.md` itself is downloaded here, not the shards, since this runs + before we know which config's data is actually needed. + """ + try: + readme_path = hf_hub_download( + repo_id=repo_id, repo_type="dataset", filename=_DATASET_CARD_FILENAME, token=token + ) + except EntryNotFoundError: + return None + + with open(readme_path, "r", encoding="utf-8") as f: + text = f.read() + + if not text.startswith("---"): + return None + front_matter_end = text.find("\n---", 3) + if front_matter_end == -1: + return None + + try: + card_metadata = yaml.safe_load(text[3:front_matter_end]) or {} + except yaml.YAMLError: + return None + + configs = card_metadata.get("configs") + if not configs: + return None + + for entry in configs: + if entry.get("default"): + return entry.get("config_name") + + return configs[0].get("config_name") + + +def _resolveConfig(repo_id: str, config: Optional[str], token: Optional[str]) -> Optional[str]: + """ + Resolves which config's files to download, without downloading any shard. + + Lists the repo's files through the Hub API (a cheap metadata call) to tell + single-config repos (files at the root) apart from multi-config repos (that + structure repeated under a folder per config), and to know each config's + name without pulling any shard data down. Returns `None` for a single-config + repo, whose files sit at the repo root. + """ + shard_index_suffix = f"/{_IMAGES_DIRNAME}/{_SHARD_INDEX_FILENAME}" + files = list_repo_files(repo_id=repo_id, repo_type="dataset", token=token) + + root_is_config = f"{_IMAGES_DIRNAME}/{_SHARD_INDEX_FILENAME}" in files + available_configs = sorted({f[: -len(shard_index_suffix)] for f in files if f.endswith(shard_index_suffix)}) + + if config is not None: + if config not in available_configs: + if not available_configs: + raise ValueError(f"'{repo_id}' has no configs, it is a single-config dataset, drop `config=`.") + raise ValueError(f"Config '{config}' was not found in '{repo_id}'. Available configs: {available_configs}.") + return config + + if root_is_config: + return None + + if not available_configs: + raise FileNotFoundError(f"No '{_IMAGES_DIRNAME}/' folder was found in '{repo_id}' or any of its subfolders.") + + default_config = _fetchDefaultConfigName(repo_id, token) + if default_config in available_configs: + return default_config + + return available_configs[0] + + def loadImageTextToTextDataset( repo_id: str, + config: Optional[str] = None, split: Optional[str] = None, cache_dir: Optional[str] = None, token: Optional[str] = None, @@ -139,13 +225,22 @@ def loadImageTextToTextDataset( """ Loads an AgML image-text-to-text shard dataset from the Hub. - This downloads `repo_id` with `snapshot_download`, builds an - `ImageTextToTextShardStore` over the downloaded `images/` shards, and loads - the metadata parquet(s) into an Arrow-backed dataset with a lazy transform - attached that decodes the `images` column into `PIL.Image` objects on - access. Splits are auto-discovered from the metadata filenames (e.g. - `train-0000-of-0001.parquet` becomes `train`); if no split can be inferred, - every row is placed under a single `train` split. + This resolves which config to load through the Hub API first, without + downloading any shard, then downloads only that config's files with + `snapshot_download`. It builds an `ImageTextToTextShardStore` over the + downloaded `images/` shards, and loads the metadata parquet(s) into an + Arrow-backed dataset with a lazy transform attached that decodes the + `images` column into `PIL.Image` objects on access. Splits are + auto-discovered from the metadata filenames (e.g. `train-0000-of-0001.parquet` + becomes `train`); if no split can be inferred, every row is placed under a + single `train` split. + + Single-config repos keep their metadata parquet(s) and `images/` folder at + the repo root, so the whole repo is downloaded either way. Multi-config + repos repeat that structure under a folder per config, in which case only + the resolved config's folder is downloaded; if `config` isn't given, the + default config named in the dataset card's `configs` YAML front matter is + used. The backing `ImageTextToTextShardStore` is attached to the returned dataset as a `.store` attribute rather than returned separately, so callers can @@ -156,6 +251,10 @@ def loadImageTextToTextDataset( ---------- repo_id : str The Hugging Face Hub dataset repo id, e.g. `Project-AgML/AgroOmni-plainziptest`. + config : str, optional + The config (subset) name to load, for repos with multiple configs each + under their own folder. If omitted, the dataset card's declared default + config is used. split : str, optional If given, only that split is returned (as a `Dataset` rather than a `DatasetDict`). cache_dir : str, optional @@ -170,15 +269,24 @@ def loadImageTextToTextDataset( `ImageTextToTextShardStore` backing the lazy image decoding attached as its `.store` attribute. """ - repo_dir = snapshot_download(repo_id=repo_id, repo_type="dataset", cache_dir=cache_dir, token=token) - images_dir = os.path.join(repo_dir, _IMAGES_DIRNAME) + resolved_config = _resolveConfig(repo_id, config, token) + allow_patterns = None if resolved_config is None else [f"{resolved_config}/*"] + repo_dir = snapshot_download( + repo_id=repo_id, + repo_type="dataset", + cache_dir=cache_dir, + token=token, + allow_patterns=allow_patterns, + ) + config_dir = repo_dir if resolved_config is None else os.path.join(repo_dir, resolved_config) + images_dir = os.path.join(config_dir, _IMAGES_DIRNAME) store = ImageTextToTextShardStore(images_dir) decode = _makeDecodeTransform(store) - split_to_files = _discoverSplits(repo_dir) + split_to_files = _discoverSplits(config_dir) dataset_dict = DatasetDict() for split_name, parquet_files in split_to_files.items(): - split_ds = Dataset.from_parquet(parquet_files) + split_ds = Dataset.from_parquet(parquet_files, split=split_name) split_ds.set_transform(decode) dataset_dict[split_name] = split_ds From 4b16b3473f9ff09319acbec68d4e4d5be45d8424 Mon Sep 17 00:00:00 2001 From: Jared Smith Date: Thu, 1 Oct 2026 19:43:17 -0700 Subject: [PATCH 3/3] some small bug fixes related to caching based on bytes and pickling of loaders --- README.md | 9 +- agml/data/image_text_to_text_loader.py | 136 ++++++++++++++++++------- 2 files changed, 106 insertions(+), 39 deletions(-) diff --git a/README.md b/README.md index 0c7fa1e6..b912e29e 100644 --- a/README.md +++ b/README.md @@ -114,7 +114,7 @@ ds = loadImageTextToTextDataset( ) ``` -`loadImageTextToTextDataset(repo_id, config=None, split=None, cache_dir=None, token=None)` returns a `DatasetDict` +`loadImageTextToTextDataset(repo_id, config=None, split=None, cache_dir=None, token=None, cache_max_bytes=...)` returns a `DatasetDict` when `split` is omitted, or a single `Dataset` when a split name is given. Splits are auto-discovered from the metadata parquet filenames (e.g. `train-0000-of-0001.parquet` becomes `train`), the same convention `load_dataset()` itself uses; if no split can be inferred, every row is placed under a single `train` split. For repos with multiple @@ -124,10 +124,11 @@ configs available in that repo. Image bytes are only read the first time a row is actually accessed (`ds["train"][0]`, a slice, or a batch), never eagerly over the whole dataset, and repeated access to the same image is cached. The `ImageTextToTextShardStore` backing that lazy decoding is attached to the returned dataset as `ds.store`, in case you want to inspect its cache, -call `ds.store.getBytes(path)` directly, or free its open shard file handles (`del ds.store`). +call `ds.store.getBytes(path)` directly, or free its open shard file handles and cached bytes (`ds.store.close()`). +The cache is bounded by total size, 256 MiB by default; pass `cache_max_bytes=` to change it, or `0` to disable it. -If you're using this with a PyTorch `DataLoader`, prefer `num_workers > 0` so each worker process opens its own -shard handles rather than sharing one across threads. +If you're using this with a PyTorch `DataLoader`, prefer `num_workers > 0`. Each worker process opens its +own shard handles (and keeps its own byte cache), whether workers are forked or spawned. ## Public Datasets diff --git a/agml/data/image_text_to_text_loader.py b/agml/data/image_text_to_text_loader.py index 7ae3e3cb..213fcc17 100644 --- a/agml/data/image_text_to_text_loader.py +++ b/agml/data/image_text_to_text_loader.py @@ -2,8 +2,7 @@ import os import re import zipfile -from collections import defaultdict -from functools import lru_cache +from collections import OrderedDict, defaultdict from typing import Optional, Union import pandas as pd @@ -43,6 +42,9 @@ # The dataset card holding the `configs` YAML front matter that names the default config. _DATASET_CARD_FILENAME = "README.md" +# The default upper bound, in bytes, on each store's cache of raw image bytes. +_DEFAULT_CACHE_MAX_BYTES = 256 * 1024 * 1024 + class ImageTextToTextShardStore: """ @@ -50,19 +52,36 @@ class ImageTextToTextShardStore: On construction, this reads `path_to_shard.parquet` into an in-memory `{path: shard_file}` index. Each shard's `zipfile.ZipFile` handle is opened - lazily on first access and kept open for reuse, and reads are cached with - an `lru_cache` since the same image path is often referenced by multiple rows. + lazily on first access and kept open for reuse. Handles are tracked per + process, so a store copied into a `DataLoader` worker (forked or pickled) + opens its own handles rather than sharing the parent's file position. + + Reads are cached in a per-store LRU bounded by total bytes (`cache_max_bytes`), + since the same image path is often referenced by multiple rows. Pass + `cache_max_bytes=0` to disable the cache. """ - def __init__(self, images_dir: str): + def __init__(self, images_dir: str, cache_max_bytes: int = _DEFAULT_CACHE_MAX_BYTES): self.images_dir = images_dir + self.cache_max_bytes = cache_max_bytes index_path = os.path.join(images_dir, _SHARD_INDEX_FILENAME) index_df = pd.read_parquet(index_path) self._path_to_shard = dict(zip(index_df["path"], index_df["shard_file"])) + self._resetHandles() + + def _resetHandles(self): + """Drops all open zip handles and cached bytes, binding the store to the current process.""" + self._pid = os.getpid() self._zip_cache = {} + self._bytes_cache = OrderedDict() + self._cached_bytes = 0 def _getZip(self, shard_file: str): - """Returns the open `ZipFile` handle for a shard, opening it on first access.""" + """Returns the open `ZipFile` handle for a shard, opening it on first access in this process.""" + if self._pid != os.getpid(): + # Forked into a new process: the inherited handles share the parent's file + # position, so abandon them (without closing, the parent still owns them). + self._resetHandles() zf = self._zip_cache.get(shard_file) if zf is None: shard_path = os.path.join(self.images_dir, shard_file) @@ -70,42 +89,77 @@ def _getZip(self, shard_file: str): self._zip_cache[shard_file] = zf return zf - @lru_cache(maxsize=4096) - def _getBytesCached(self, path: str) -> bytes: - """Reads raw image bytes for a path, caching repeat reads of the same path.""" - shard_file = self._path_to_shard[path] - return self._getZip(shard_file).read(path) - def getBytes(self, path: str) -> bytes: - """Public entry point for reading the raw bytes of a single image path.""" - return self._getBytesCached(path) + """Reads the raw bytes of a single image path, caching repeat reads of the same path.""" + if self._pid != os.getpid(): + self._resetHandles() + raw_bytes = self._bytes_cache.get(path) + if raw_bytes is not None: + self._bytes_cache.move_to_end(path) + return raw_bytes + + shard_file = self._path_to_shard[path] + raw_bytes = self._getZip(shard_file).read(path) + + if len(raw_bytes) <= self.cache_max_bytes: + self._bytes_cache[path] = raw_bytes + self._cached_bytes += len(raw_bytes) + while self._cached_bytes > self.cache_max_bytes: + _, evicted = self._bytes_cache.popitem(last=False) + self._cached_bytes -= len(evicted) + return raw_bytes + + def close(self): + """Closes every open shard handle and clears the byte cache. Handles reopen lazily on next access.""" + if getattr(self, "_pid", None) == os.getpid(): + for zf in self._zip_cache.values(): + try: + zf.close() + except Exception: + pass + self._resetHandles() + + def __getstate__(self): + # Open file handles can't be pickled (e.g. into spawned `DataLoader` workers), + # and the byte cache shouldn't be shipped along, so drop both. + state = self.__dict__.copy() + state["_zip_cache"] = {} + state["_bytes_cache"] = OrderedDict() + state["_cached_bytes"] = 0 + state["_pid"] = None + return state def __del__(self): # Best effort cleanup, the process may already be tearing down. - for zf in getattr(self, "_zip_cache", {}).values(): - try: - zf.close() - except Exception: - pass + try: + self.close() + except Exception: + pass + + +class _DecodeTransform: + """ + The lazy `set_transform` callback that decodes the `images` column via `store`. + This is a module-level class rather than a closure so the dataset it is + attached to can be pickled into spawned `DataLoader` workers. + """ -def _makeDecodeTransform(store: ImageTextToTextShardStore): - """Builds the lazy `set_transform` callback that decodes the `images` column via `store`.""" + def __init__(self, store: ImageTextToTextShardStore): + self.store = store - def decode(batch): + def __call__(self, batch): decoded_rows = [] for images in batch["images"]: decoded_row = [] for image_ref in images: path = image_ref["path"] if isinstance(image_ref, dict) else image_ref - raw_bytes = store.getBytes(path) + raw_bytes = self.store.getBytes(path) decoded_row.append(Image.open(io.BytesIO(raw_bytes))) decoded_rows.append(decoded_row) batch["images"] = decoded_rows return batch - return decode - def _discoverSplits(config_dir: str) -> dict: """ @@ -135,7 +189,7 @@ def _discoverSplits(config_dir: str) -> dict: return split_to_files -def _fetchDefaultConfigName(repo_id: str, token: Optional[str]) -> Optional[str]: +def _fetchDefaultConfigName(repo_id: str, token: Optional[str], cache_dir: Optional[str]) -> Optional[str]: """ Reads the default config name out of the dataset card's `configs` YAML front matter. @@ -149,7 +203,11 @@ def _fetchDefaultConfigName(repo_id: str, token: Optional[str]) -> Optional[str] """ try: readme_path = hf_hub_download( - repo_id=repo_id, repo_type="dataset", filename=_DATASET_CARD_FILENAME, token=token + repo_id=repo_id, + repo_type="dataset", + filename=_DATASET_CARD_FILENAME, + token=token, + cache_dir=cache_dir, ) except EntryNotFoundError: return None @@ -179,7 +237,9 @@ def _fetchDefaultConfigName(repo_id: str, token: Optional[str]) -> Optional[str] return configs[0].get("config_name") -def _resolveConfig(repo_id: str, config: Optional[str], token: Optional[str]) -> Optional[str]: +def _resolveConfig( + repo_id: str, config: Optional[str], token: Optional[str], cache_dir: Optional[str] +) -> Optional[str]: """ Resolves which config's files to download, without downloading any shard. @@ -208,7 +268,7 @@ def _resolveConfig(repo_id: str, config: Optional[str], token: Optional[str]) -> if not available_configs: raise FileNotFoundError(f"No '{_IMAGES_DIRNAME}/' folder was found in '{repo_id}' or any of its subfolders.") - default_config = _fetchDefaultConfigName(repo_id, token) + default_config = _fetchDefaultConfigName(repo_id, token, cache_dir) if default_config in available_configs: return default_config @@ -221,6 +281,7 @@ def loadImageTextToTextDataset( split: Optional[str] = None, cache_dir: Optional[str] = None, token: Optional[str] = None, + cache_max_bytes: int = _DEFAULT_CACHE_MAX_BYTES, ) -> Union["DatasetDict", "Dataset"]: """ Loads an AgML image-text-to-text shard dataset from the Hub. @@ -244,7 +305,7 @@ def loadImageTextToTextDataset( The backing `ImageTextToTextShardStore` is attached to the returned dataset as a `.store` attribute rather than returned separately, so callers can - inspect its cache or free its open file handles (`del ds.store`) without + inspect its cache or free its open file handles (`ds.store.close()`) without changing the common `ds = loadImageTextToTextDataset(...)` call shape. Parameters @@ -258,9 +319,14 @@ def loadImageTextToTextDataset( split : str, optional If given, only that split is returned (as a `Dataset` rather than a `DatasetDict`). cache_dir : str, optional - The local directory to cache the downloaded repo in. + The local directory to cache the downloaded repo and the Arrow + files built from its metadata parquet(s) in. token : str, optional A Hugging Face Hub token, needed for private or gated repos. + cache_max_bytes : int, optional + The upper bound, in bytes, on the store's in-memory cache of raw image + bytes (256 MiB by default). Each `DataLoader` worker keeps its own + cache of up to this size. Pass `0` to disable caching. Returns ------- @@ -269,7 +335,7 @@ def loadImageTextToTextDataset( `ImageTextToTextShardStore` backing the lazy image decoding attached as its `.store` attribute. """ - resolved_config = _resolveConfig(repo_id, config, token) + resolved_config = _resolveConfig(repo_id, config, token, cache_dir) allow_patterns = None if resolved_config is None else [f"{resolved_config}/*"] repo_dir = snapshot_download( repo_id=repo_id, @@ -280,13 +346,13 @@ def loadImageTextToTextDataset( ) config_dir = repo_dir if resolved_config is None else os.path.join(repo_dir, resolved_config) images_dir = os.path.join(config_dir, _IMAGES_DIRNAME) - store = ImageTextToTextShardStore(images_dir) - decode = _makeDecodeTransform(store) + store = ImageTextToTextShardStore(images_dir, cache_max_bytes=cache_max_bytes) + decode = _DecodeTransform(store) split_to_files = _discoverSplits(config_dir) dataset_dict = DatasetDict() for split_name, parquet_files in split_to_files.items(): - split_ds = Dataset.from_parquet(parquet_files, split=split_name) + split_ds = Dataset.from_parquet(parquet_files, split=split_name, cache_dir=cache_dir) split_ds.set_transform(decode) dataset_dict[split_name] = split_ds