diff --git a/src/gasbench/dataset/config.py b/src/gasbench/dataset/config.py index 84c197e..b7d1e07 100644 --- a/src/gasbench/dataset/config.py +++ b/src/gasbench/dataset/config.py @@ -78,6 +78,7 @@ class BenchmarkDatasetConfig: # Examples: filter_column: "label", filter_value: "fake" or 0 filter_column: Optional[str] = None filter_value: Optional[Union[str, int, float, bool]] = None + filter_values: Optional[Dict[str, Union[str, int, float, bool]]] = None # For holdout datasets: stores the original name before obfuscation original_name: Optional[str] = None @@ -357,6 +358,14 @@ def validate_dataset_config( or not prefixes or any(not isinstance(p, str) or not p for p in prefixes)): errors.append(f"Dataset '{dataset_name}': s3_prefixes requires S3 and nonempty key prefixes") + filter_values = config_dict.get("filter_values") + if filter_values is not None and ( + not isinstance(filter_values, dict) + or not filter_values + or any(not isinstance(key, str) or not key or value is None for key, value in filter_values.items()) + ): + errors.append(f"Dataset '{dataset_name}': filter_values must map nonempty column names to values") + numeric_fields = [ "media_per_archive", "archives_per_dataset", @@ -406,6 +415,7 @@ def _dataset_dict_to_config(d: dict, **overrides) -> BenchmarkDatasetConfig: "notes": d.get("notes"), "filter_column": d.get("filter_column"), "filter_value": d.get("filter_value"), + "filter_values": d.get("filter_values"), "generator_family": d.get("generator_family"), "content_category": d.get("content_category"), } diff --git a/src/gasbench/dataset/configs/real_audio.yaml b/src/gasbench/dataset/configs/real_audio.yaml index 4d6f44f..515ab60 100644 --- a/src/gasbench/dataset/configs/real_audio.yaml +++ b/src/gasbench/dataset/configs/real_audio.yaml @@ -896,3 +896,60 @@ datasets: archives_per_dataset: -1 notes: 'Released from v24 bmcore holdouts; exact local subset. ' hf_revision: 0699bb1d85463a1cb5928db83bc61bdd127b4590 + + # Released v26 holdouts — pinned upstream HF sources. + - name: v26-real-galgame-speech-asr-16khz + path: litagin/Galgame_Speech_ASR_16kHz + hf_revision: 3fb86654222b3f0af0f7c332ae6a0ef9752a9451 + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: tar + include_paths: + - data/galgame-speech-asr-16kHz-train- + modality: audio + media_type: real + - name: v26-real-sada2022-arabic-tts + path: mohamedalgmaal/sada2022-arabic-tts + hf_revision: 10870c1a64d17a24a791ee08e8297f216f96a447 + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: wav + include_paths: + - batch_ + modality: audio + media_type: real + - name: v26-real-starrail-voice + path: simon3000/starrail-voice + hf_revision: 1e5f1e8658ab2e618c4815cf64ddf905c361067d + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: parquet + include_paths: + - data/train- + data_columns: + - audio + modality: audio + media_type: real + - name: v26-real-worldspeech + path: disco-eth/WorldSpeech + hf_revision: b76b85395039f50332fbc421f7c098ee2f2653d7 + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: parquet + include_paths: + - data/af_za/test- + - data/am_et/test- + - data/ar_bh/test- + - data/ar_dz/test- + - data/ar_eg/test- + - data/ar_iq/test- + - data/ar_kw/test- + - data/ar_ma/test- + - data/ar_sa/test- + - data/ar_tn/test- + - data/ar_un/test- + - data/az_az/test- + data_columns: + - audio + modality: audio + media_type: real diff --git a/src/gasbench/dataset/configs/real_images.yaml b/src/gasbench/dataset/configs/real_images.yaml index 28c0027..a4acd69 100644 --- a/src/gasbench/dataset/configs/real_images.yaml +++ b/src/gasbench/dataset/configs/real_images.yaml @@ -1288,3 +1288,86 @@ datasets: archives_per_dataset: -1 notes: 'Released from v24 bmcore holdouts; exact local subset. ' hf_revision: a3cf2557e18a832c4b2921c8d20948202eab54f5 + + # Released v26 holdouts — pinned upstream HF sources. + - name: v25-real-beans-test + path: AI-Lab-Makerere/beans + hf_revision: 27aa014ce09b193e1a6f58112d4a66e0eddb69c5 + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: parquet + include_paths: + - data/test- + data_columns: + - image + modality: image + media_type: real + - name: v25-real-beans-train + path: AI-Lab-Makerere/beans + hf_revision: 27aa014ce09b193e1a6f58112d4a66e0eddb69c5 + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: parquet + include_paths: + - data/train- + data_columns: + - image + modality: image + media_type: real + - name: v25-real-beans-validation + path: AI-Lab-Makerere/beans + hf_revision: 27aa014ce09b193e1a6f58112d4a66e0eddb69c5 + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: parquet + include_paths: + - data/validation- + data_columns: + - image + modality: image + media_type: real + - name: v25-real-oxford-iiit-pet-test + path: timm/oxford-iiit-pet + hf_revision: 089695c834a7deb60505b7cc506672db1c31a6aa + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: parquet + include_paths: + - data/test- + data_columns: + - image + modality: image + media_type: real + - name: v25-real-oxford-iiit-pet-train + path: timm/oxford-iiit-pet + hf_revision: 089695c834a7deb60505b7cc506672db1c31a6aa + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: parquet + include_paths: + - data/train- + data_columns: + - image + modality: image + media_type: real + - name: v26-real-dcvlm-pool-large-dvqa + path: mlfoundations/dcvlm_pool_large + hf_revision: 479162ae38203029e06807038cd2e5e545bfd2e0 + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: tar + include_paths: + - dvqa/ + modality: image + media_type: real + - name: v26-real-puffin-16m + path: ACERobotics/Puffin-16M + hf_revision: 2112b3502f9ea0cc1d31d2c5a2b7f2259b52cffd + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: zip + include_paths: + - Puffin-Cam-15M/ + - Puffin-Traj-1M/ + modality: image + media_type: real diff --git a/src/gasbench/dataset/configs/real_videos.yaml b/src/gasbench/dataset/configs/real_videos.yaml index a925063..7ed80f2 100644 --- a/src/gasbench/dataset/configs/real_videos.yaml +++ b/src/gasbench/dataset/configs/real_videos.yaml @@ -1339,3 +1339,155 @@ datasets: archives_per_dataset: -1 notes: 'Released from v24 bmcore holdouts; exact local subset. ' hf_revision: f0af08aece50ff99a9c399f063d71e455d610bbf + + # Released v26 holdouts — pinned upstream HF sources. + - name: v25-rendered-physical-interaction-bowling-00000 + path: nvidia/PhysicalAI-WorldModel-Synthetic-Physical-Interaction-Scenes + hf_revision: bdafca2a9c575629e47509f547e27ab7f3680369 + media_per_archive: 250 + archives_per_dataset: 1 + source_format: tar + include_paths: + - videos/bowling/videos-bowling-00000.tar + archive_include_paths: + - Bowler.mp4 + - Overhead.mp4 + - Pins.mp4 + - Side.mp4 + modality: video + media_type: real + - name: v25-rendered-physical-interaction-bowling-00042 + path: nvidia/PhysicalAI-WorldModel-Synthetic-Physical-Interaction-Scenes + hf_revision: bdafca2a9c575629e47509f547e27ab7f3680369 + media_per_archive: 250 + archives_per_dataset: 1 + source_format: tar + include_paths: + - videos/bowling/videos-bowling-00042.tar + archive_include_paths: + - Bowler.mp4 + - Overhead.mp4 + - Pins.mp4 + - Side.mp4 + modality: video + media_type: real + - name: v25-rendered-physical-interaction-bowling-00063 + path: nvidia/PhysicalAI-WorldModel-Synthetic-Physical-Interaction-Scenes + hf_revision: bdafca2a9c575629e47509f547e27ab7f3680369 + media_per_archive: 250 + archives_per_dataset: 1 + source_format: tar + include_paths: + - videos/bowling/videos-bowling-00063.tar + archive_include_paths: + - Bowler.mp4 + - Overhead.mp4 + - Pins.mp4 + - Side.mp4 + modality: video + media_type: real + - name: v25-rendered-physical-interaction-bowling-00084 + path: nvidia/PhysicalAI-WorldModel-Synthetic-Physical-Interaction-Scenes + hf_revision: bdafca2a9c575629e47509f547e27ab7f3680369 + media_per_archive: 250 + archives_per_dataset: 1 + source_format: tar + include_paths: + - videos/bowling/videos-bowling-00084.tar + archive_include_paths: + - Bowler.mp4 + - Overhead.mp4 + - Pins.mp4 + - Side.mp4 + modality: video + media_type: real + - name: v25-rendered-physicalai-worldmodel-synthetic-warehouse-operations-scenes-artifacts-forklift-human-nearmiss-nearmiss-artifacts-00005 + path: nvidia/PhysicalAI-WorldModel-Synthetic-Warehouse-Operations-Scenes + hf_revision: d5b88d3abcf659f304a107f4336b71b4e2159133 + media_per_archive: 250 + archives_per_dataset: 1 + source_format: tar + include_paths: + - artifacts/forklift_human_nearmiss/nearmiss-artifacts-00005.tar + archive_include_paths: + - .rgb.mp4 + modality: video + media_type: real + - name: v25-rendered-physicalai-worldmodel-synthetic-warehouse-operations-scenes-artifacts-forklift-human-nearmiss-nearmiss-artifacts-00016 + path: nvidia/PhysicalAI-WorldModel-Synthetic-Warehouse-Operations-Scenes + hf_revision: d5b88d3abcf659f304a107f4336b71b4e2159133 + media_per_archive: 250 + archives_per_dataset: 1 + source_format: tar + include_paths: + - artifacts/forklift_human_nearmiss/nearmiss-artifacts-00016.tar + archive_include_paths: + - .rgb.mp4 + modality: video + media_type: real + - name: v25-rendered-physicalai-worldmodel-synthetic-warehouse-operations-scenes-artifacts-forklift-human-nearmiss-nearmiss-artifacts-00028 + path: nvidia/PhysicalAI-WorldModel-Synthetic-Warehouse-Operations-Scenes + hf_revision: d5b88d3abcf659f304a107f4336b71b4e2159133 + media_per_archive: 250 + archives_per_dataset: 1 + source_format: tar + include_paths: + - artifacts/forklift_human_nearmiss/nearmiss-artifacts-00028.tar + archive_include_paths: + - .rgb.mp4 + modality: video + media_type: real + - name: v25-rendered-physicalai-worldmodel-synthetic-warehouse-operations-scenes-artifacts-forklift-human-nearmiss-nearmiss-artifacts-00039 + path: nvidia/PhysicalAI-WorldModel-Synthetic-Warehouse-Operations-Scenes + hf_revision: d5b88d3abcf659f304a107f4336b71b4e2159133 + media_per_archive: 250 + archives_per_dataset: 1 + source_format: tar + include_paths: + - artifacts/forklift_human_nearmiss/nearmiss-artifacts-00039.tar + archive_include_paths: + - .rgb.mp4 + modality: video + media_type: real + - name: v25-rendered-physicalai-worldmodel-synthetic-warehouse-operations-scenes-artifacts-forklift-human-nearmiss-nearmiss-artifacts-00050 + path: nvidia/PhysicalAI-WorldModel-Synthetic-Warehouse-Operations-Scenes + hf_revision: d5b88d3abcf659f304a107f4336b71b4e2159133 + media_per_archive: 250 + archives_per_dataset: 1 + source_format: tar + include_paths: + - artifacts/forklift_human_nearmiss/nearmiss-artifacts-00050.tar + archive_include_paths: + - .rgb.mp4 + modality: video + media_type: real + - name: v26-real-physicalai-robotics-manipulation-objects + path: nvidia/PhysicalAI-Robotics-Manipulation-Objects + hf_revision: 538b6833148b4731e4ea30fe181cc4eecfc13440 + media_per_archive: 250 + archives_per_dataset: 1 + source_format: mp4 + include_paths: + - .rgb/ + modality: video + media_type: real + - name: v26-real-pick-tblock-mp-centered-0701-rendered + path: younghyopark/pick_tblock_mp_centered_0701_rendered + hf_revision: b353fb987b6cc55524fb144ea61632168bc2ca5f + media_per_archive: 250 + archives_per_dataset: 1 + source_format: mp4 + include_paths: + - videos/chunk-000/ + modality: video + media_type: real + - name: v26-real-vidgen-1m + path: Fudan-FUXI/VIDGEN-1M + hf_revision: 9158f218b6b81694a7809b6ee74e30003d58a70e + media_per_archive: 250 + archives_per_dataset: 1 + source_format: zip + include_paths: + - VidGen_video_ + modality: video + media_type: real diff --git a/src/gasbench/dataset/configs/synthetic_audio.yaml b/src/gasbench/dataset/configs/synthetic_audio.yaml index ae8c069..dd5823c 100644 --- a/src/gasbench/dataset/configs/synthetic_audio.yaml +++ b/src/gasbench/dataset/configs/synthetic_audio.yaml @@ -900,3 +900,89 @@ datasets: archives_per_dataset: -1 notes: 'Released from v24 bmcore holdouts; exact local subset. ' hf_revision: 04188209b3135617c0c682c432c2168e01b35011 + + # Released v26 holdouts — pinned upstream HF sources. + - name: v25-Fake-Synthetic-Turkish-TTS-Data + path: Anilosan15/Synthetic_Turkish_TTS_Data + hf_revision: f1d0e40efb27b5e78cdd6dbd6f6b601ab2ce9ae0 + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: parquet + include_paths: + - data/train- + data_columns: + - audio + modality: audio + media_type: synthetic + - name: v25-Fake-qwen-clones-4m-en-audio-train-00000-of-00080 + path: SynDataLab-EN/qwen-clones-4m-en + hf_revision: 38769c7d96e5767035555a66ce57b1076a7dccf1 + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: parquet + include_paths: + - data/train-00000-of-00080.parquet + data_columns: + - audio + modality: audio + media_type: synthetic + - name: v25-Fake-qwen-clones-4m-en-audio-train-00035-of-00080 + path: SynDataLab-EN/qwen-clones-4m-en + hf_revision: 38769c7d96e5767035555a66ce57b1076a7dccf1 + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: parquet + include_paths: + - data/train-00035-of-00080.parquet + data_columns: + - audio + modality: audio + media_type: synthetic + - name: v25-Fake-qwen-clones-4m-en-audio-train-00052-of-00080 + path: SynDataLab-EN/qwen-clones-4m-en + hf_revision: 38769c7d96e5767035555a66ce57b1076a7dccf1 + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: parquet + include_paths: + - data/train-00052-of-00080.parquet + data_columns: + - audio + modality: audio + media_type: synthetic + - name: v25-Fake-qwen-clones-4m-en-audio-train-00074-of-00080 + path: SynDataLab-EN/qwen-clones-4m-en + hf_revision: 38769c7d96e5767035555a66ce57b1076a7dccf1 + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: parquet + include_paths: + - data/train-00074-of-00080.parquet + data_columns: + - audio + modality: audio + media_type: synthetic + - name: v26-synthetic-echo-clones-4m-en + path: SynDataLab-EN/echo-clones-4m-en + hf_revision: aebc7c787c64e918b0130c04d3878f17a602c1df + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: parquet + include_paths: + - data/train- + data_columns: + - audio + modality: audio + media_type: synthetic + - name: v26-synthetic-steak + path: DEEL-AI/STEAK + hf_revision: 4ee840a09e582d23d2ce88b7f99a1996ffe9aec6 + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: parquet + include_paths: + - data/STEAK- + data_columns: + - audio + modality: audio + media_type: synthetic diff --git a/src/gasbench/dataset/configs/synthetic_images.yaml b/src/gasbench/dataset/configs/synthetic_images.yaml index cf0cabb..5c1d9b0 100644 --- a/src/gasbench/dataset/configs/synthetic_images.yaml +++ b/src/gasbench/dataset/configs/synthetic_images.yaml @@ -1173,3 +1173,153 @@ datasets: archives_per_dataset: -1 notes: 'Released from v24 bmcore holdouts; exact local subset. ' hf_revision: 5649a33952bb8faabf3b861ad3b5024374509e72 + + # Released v26 holdouts — pinned upstream HF sources. + - name: v25-synthetic-10D-500K-images-Animals-bagel + path: addisonwu05/10D-500K + hf_revision: d5e2472237a1dc612d9b1d1bfaed8b9b90fbc0e1 + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: parquet + include_paths: + - data/train- + data_columns: + - image + filter_values: + domain: Animals + model: bagel + modality: image + media_type: synthetic + - name: v25-synthetic-10D-500K-images-Buildings-emu + path: addisonwu05/10D-500K + hf_revision: d5e2472237a1dc612d9b1d1bfaed8b9b90fbc0e1 + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: parquet + include_paths: + - data/train- + data_columns: + - image + filter_values: + domain: Buildings + model: emu + modality: image + media_type: synthetic + - name: v25-synthetic-10D-500K-images-Food_and_Drinks-janus + path: addisonwu05/10D-500K + hf_revision: d5e2472237a1dc612d9b1d1bfaed8b9b90fbc0e1 + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: parquet + include_paths: + - data/train- + data_columns: + - image + filter_values: + domain: Food_and_Drinks + model: janus + modality: image + media_type: synthetic + - name: v25-synthetic-10D-500K-images-Interior_Spaces-mmada + path: addisonwu05/10D-500K + hf_revision: d5e2472237a1dc612d9b1d1bfaed8b9b90fbc0e1 + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: parquet + include_paths: + - data/train- + data_columns: + - image + filter_values: + domain: Interior_Spaces + model: mmada + modality: image + media_type: synthetic + - name: v25-synthetic-10D-500K-images-People-showo2 + path: addisonwu05/10D-500K + hf_revision: d5e2472237a1dc612d9b1d1bfaed8b9b90fbc0e1 + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: parquet + include_paths: + - data/train- + data_columns: + - image + filter_values: + domain: People + model: showo2 + modality: image + media_type: synthetic + - name: v25-synthetic-Persona-Fluxed-10k-2608_brazil-train + path: retowyss/Persona-Fluxed-10k-2608 + hf_revision: abd1938c598a1f67a242bb35cd9740a5e1870e78 + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: png + include_paths: + - brazil/train/ + modality: image + media_type: synthetic + - name: v25-synthetic-Persona-Fluxed-10k-2608_france-train + path: retowyss/Persona-Fluxed-10k-2608 + hf_revision: abd1938c598a1f67a242bb35cd9740a5e1870e78 + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: png + include_paths: + - france/train/ + modality: image + media_type: synthetic + - name: v25-synthetic-Persona-Fluxed-10k-2608_japan-train + path: retowyss/Persona-Fluxed-10k-2608 + hf_revision: abd1938c598a1f67a242bb35cd9740a5e1870e78 + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: png + include_paths: + - japan/train/ + modality: image + media_type: synthetic + - name: v25-synthetic-Syn-Vis-v0-images-headshot + path: retowyss/Syn-Vis-v0 + hf_revision: 100262732989e77f38cd831d70a376a93735006a + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: png + include_paths: + - images/headshot/ + modality: image + media_type: synthetic + - name: v26-synthetic-per-400k + path: a1557811266/Per-400k + hf_revision: bbb01675182feca1434ae3fbb97a9355d9337a25 + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: zip + include_paths: + - person_data/img/ + modality: image + media_type: synthetic + - name: v26-synthetic-sora-explore-dataset + path: abrar71/sora-explore-dataset + hf_revision: 79559046d3782bb69147209bbbd72a787a326452 + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: parquet + include_paths: + - dataset/train- + data_columns: + - image + modality: image + media_type: synthetic + - name: v26-synthetic-stylegan-face + path: gojay/StyleGAN-Face + hf_revision: d7592242ec59376dd1c418dd0353c2361631ac6e + media_per_archive: 1000 + archives_per_dataset: 1 + source_format: zip + include_paths: + - age_ + - smile_ + modality: image + media_type: synthetic diff --git a/src/gasbench/dataset/configs/synthetic_videos.yaml b/src/gasbench/dataset/configs/synthetic_videos.yaml index e6c5b6d..626b709 100644 --- a/src/gasbench/dataset/configs/synthetic_videos.yaml +++ b/src/gasbench/dataset/configs/synthetic_videos.yaml @@ -1297,3 +1297,42 @@ datasets: archives_per_dataset: -1 notes: 'Released from v24 bmcore holdouts; exact local subset. ' hf_revision: fa7713f57249caf009200ab0f875c07cd3e20986 + + # Released v26 holdouts — pinned upstream HF sources. + - name: v26-semisynthetic-vera-layered-video-dataset + path: netflix/Vera-Layered-Video-Dataset + hf_revision: 8e0b98ee9bce66fdae345e75aa766ef7c0a04d4e + media_per_archive: 250 + archives_per_dataset: 1 + source_format: mp4 + include_paths: + - train-set/49-frames/synthetic-obj-add/composite/ + - train-set/81-frames/synthetic-obj-add/composite/ + modality: video + media_type: semisynthetic + - name: v26-synthetic-text-2-video-human-preferences-genmo-mochi-1 + path: Rapidata/text-2-video-human-preferences-genmo-mochi-1 + hf_revision: 9b8c6dbba6ba4e034adaa509550e53d81e3b7148 + media_per_archive: 250 + archives_per_dataset: 1 + source_format: mp4 + include_paths: + - Videos/ + modality: video + media_type: synthetic + - name: v26-synthetic-vbench-2-0-sampled-videos + path: Vchitect/VBench-2.0_sampled_videos + hf_revision: f9ee15efcfb86d96bd5379e8d85f8d69b901aad1 + media_per_archive: 250 + archives_per_dataset: 1 + source_format: mp4 + modality: video + media_type: synthetic + - name: v26-synthetic-vbench-i2v-sampled-video + path: Vchitect/VBench-I2V_sampled_video + hf_revision: d1977eff951fad22ac7ac11059f6939f570a78eb + media_per_archive: 250 + archives_per_dataset: 1 + source_format: mp4 + modality: video + media_type: synthetic diff --git a/src/gasbench/dataset/download/core.py b/src/gasbench/dataset/download/core.py index 9bc3115..84b23b7 100644 --- a/src/gasbench/dataset/download/core.py +++ b/src/gasbench/dataset/download/core.py @@ -141,11 +141,14 @@ def download_and_extract( max_files_to_list = n_files if n_files > 0 else None # ── Filtered sequential mode ───────────────────────────────────── - # When filter_column/filter_value are set, download parquet shards + # When row filters are set, download parquet shards # one at a time and stop as soon as the sample target is hit. filter_column = getattr(dataset, "filter_column", None) filter_value = getattr(dataset, "filter_value", None) + filter_values = dict(getattr(dataset, "filter_values", None) or {}) if filter_column and filter_value is not None: + filter_values[filter_column] = filter_value + if filter_values: try: all_filenames = _list_remote_dataset_files( dataset.path, @@ -172,7 +175,7 @@ def download_and_extract( target_total = media_per_archive * archives_per_dataset logger.info( - f"Filtered mode: {dataset.name} — {filter_column}=={filter_value}, " + f"Filtered mode: {dataset.name} — {filter_values}, " f"target {target_total} samples from {len(all_filenames)} shards" ) yield from _download_filtered_sequential( diff --git a/src/gasbench/dataset/download/extract.py b/src/gasbench/dataset/download/extract.py index 318000e..9b89269 100644 --- a/src/gasbench/dataset/download/extract.py +++ b/src/gasbench/dataset/download/extract.py @@ -426,15 +426,19 @@ def _process_parquet( filter_column = getattr(dataset, "filter_column", None) filter_value = getattr(dataset, "filter_value", None) + filter_values = dict(getattr(dataset, "filter_values", None) or {}) if filter_column and filter_value is not None: - if filter_column not in df.columns: - logger.warning(f"filter_column '{filter_column}' not in {source_path.name}; skipping") + filter_values[filter_column] = filter_value + for column, value in filter_values.items(): + if column not in df.columns: + logger.warning(f"filter column '{column}' not in {source_path.name}; skipping") return - df = df[df[filter_column] == filter_value] + df = df[df[column] == value] if df.empty: - logger.info(f"No rows matching {filter_column}=={filter_value} in {source_path.name}") + logger.info(f"No rows matching {filter_values} in {source_path.name}") return - logger.info(f" {source_path.name}: {len(df)} rows after filter ({filter_column}=={filter_value})") + if filter_values: + logger.info(f" {source_path.name}: {len(df)} rows after filter ({filter_values})") # Frames-parquet: video dataset stored as one row per frame, grouped by video_id if dataset.modality == "video" and "video_id" in df.columns: @@ -689,4 +693,3 @@ def extract_number(filepath): except Exception as e: logger.warning(f"Error processing frame directory {frame_dir}: {e}") return - diff --git a/tests/integration/test_parquet_processing.py b/tests/integration/test_parquet_processing.py index 6054863..14dae8e 100644 --- a/tests/integration/test_parquet_processing.py +++ b/tests/integration/test_parquet_processing.py @@ -144,6 +144,41 @@ def download(paths, root, **kwargs): assert all(sample["label"] == label and sample["media_type"] == media_type for sample in samples) +def test_multiple_yaml_filters_collect_only_matching_rows_across_shards(tmp_path, monkeypatch): + config_path = tmp_path / "datasets.yaml" + config_path.write_text(yaml.safe_dump({"datasets": [{ + "name": "domain-model", "path": "test/images", "modality": "image", + "media_type": "synthetic", "source_format": "parquet", + "filter_values": {"domain": "Animals", "model": "bagel"}, + }]})) + (config,) = load_datasets_from_yaml(str(config_path))["image"] + filenames = [f"shard-{i}.parquet" for i in range(4)] + monkeypatch.setattr(core, "_list_remote_dataset_files", lambda *a, **k: filenames) + monkeypatch.setattr(core, "_get_download_urls", lambda path, files, *a: list(files)) + downloaded = [] + + def download(paths, root, **kwargs): + for filename in paths: + path = root / filename + pd.DataFrame([ + {"image": png(10), "domain": "Animals", "model": "bagel"}, + {"image": png(20), "domain": "Animals", "model": "emu"}, + {"image": png(30), "domain": "Buildings", "model": "bagel"}, + ]).to_parquet(path) + downloaded.append(filename) + yield path + + monkeypatch.setattr(core, "_stream_downloads", download) + monkeypatch.setattr(core, "download_files", lambda *a, **k: list(download(*a, **k))) + samples = list(core.download_and_extract( + config, media_per_archive=3, archives_per_dataset=1, + temp_dir=str(tmp_path), force_download=True, seed=7, + )) + + assert len(samples) == len(downloaded) == 3 + assert all(sample["domain"] == "Animals" and sample["model"] == "bagel" for sample in samples) + + @pytest.mark.parametrize("metadata", [{}, {"label": 1}]) def test_zero_filter_yields_nothing_when_label_is_missing_or_unmatched(parquet, metadata): config = BenchmarkDatasetConfig(