[1/3] audioprotopnet: data layer updates by nkundiushuti · Pull Request #186 · earthspecies/avex

nkundiushuti · 2026-05-29T14:52:21Z

Closes #188.

Related to #184 (audioprotopnet model PR). Splits out data-layer changes so they can be reviewed and merged independently.

Replace sklearn train_test_split with esp_data PandasBackend in TrainValSplitTransform — removes sklearn dependency, supports non-DataFrame backends
Add worker_timeout param to build_dataloaders
Add birdset_train_splits module; registers birdset_train dataset name used in benchmark configs
Split ebird_taxonomy.json into v2021/v2025; load() now requires explicit version arg
Add fill_labels_from_answer transform to beans benchmark configs
Rename birdset → birdset_train in birdset benchmark configs
Add slurm config variants for beans detection subsets
Exclude avex/data/ from large-file, ruff, and codespell hooks (taxonomy JSON files contain Latin species names)

Dependencies

None — merges directly into main.

Test plan

pytest tests/unittests/test_ebird_taxonomy.py
pytest tests/unittests/ (full suite)

- Replace sklearn train_test_split with esp_data PandasBackend in TrainValSplitTransform - Add worker_timeout param to build_dataloaders - Add birdset_train_splits module (registers birdset_train dataset) - Split ebird_taxonomy.json into v2021/v2025; load() now requires version arg - Add fill_labels_from_answer transform to beans benchmark configs - Rename birdset -> birdset_train in birdset benchmark configs - Add slurm benchmark config variants for beans

Copilot

Pull request overview

This PR updates AVEX’s data layer to support the AudioProtoPNet integration by removing the sklearn dependency from dataset splitting, introducing versioned eBird taxonomies, and adjusting benchmark configs (BirdSet + BEANS) to match upstream dataset changes and repair label issues.

Changes:

Replace sklearn-based train/val splitting with an esp_data-compatible implementation and add a BEANS label-repair transform.
Add a temporary birdset_train dataset registration to restore legacy BirdSet train splits used by existing benchmark configs.
Version the bundled eBird taxonomy (v2021 / v2025) and add tests + config updates; add DataLoader worker_timeout.

Reviewed changes

Copilot reviewed 13 out of 15 changed files in this pull request and generated 22 comments.

Show a summary per file

File	Description
`avex/data/transforms.py`	Removes sklearn split usage; updates transforms to work with `esp_data` backends; adds `fill_labels_from_answer`.
`avex/data/dataset.py`	Imports BirdSet registration module and adds `worker_timeout` plumbed into DataLoaders.
`avex/data/birdset_train_splits.py`	Registers `birdset_train` dataset with legacy split paths to keep BirdSet benchmarks working.
`avex/data/ebird_taxonomy.py`	Adds versioned taxonomy loader requiring an explicit version argument.
`tests/unittests/test_ebird_taxonomy.py`	Adds tests for versioned taxonomy loading and bundled JSON presence.
`pytest.ini`	Adds an `integration` marker description.
`configs/data_configs/benchmark_birdset.yml`	Renames dataset usage from `birdset` to `birdset_train`.
`configs/data_configs/benchmark_beans.yml`	Adds `fill_labels_from_answer` to relevant BEANS detection transforms.
`configs/data_configs/benchmark_beans_slurm.yml`	Adds a SLURM-oriented BEANS benchmark config variant.
`configs/data_configs/benchmark_beans_remaining_slurm.yml`	Adds a “remaining tasks” SLURM config for BEANS subsets.
`configs/data_configs/benchmark_beans_detection_slurm.yml`	Adds a detection-only SLURM config variant.
`configs/data_configs/benchmark_beans_detection_remaining_slurm.yml`	Adds a “remaining detection tasks” SLURM config variant.
`.pre-commit-config.yaml`	Excludes `avex/data/` from several hooks (large files, ruff, codespell).

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

- transforms: add DataFrame path using iloc/unique for pandas compatibility; backend API path retained for esp_data Polars/other backends - pre-commit: narrow large-file and codespell excludes to ebird_taxonomy_v*.json; remove ruff/ruff-format avex/data/ exclude (ruff only targets Python anyway) - ebird_taxonomy: remove unused logger; fix FileNotFoundError docstring indent - birdset_train_splits: add missing Returns/Raises/Yields docstring sections - dataset: document worker_timeout param in build_dataloaders docstring - test_ebird_taxonomy: annotate version param as EbirdTaxonomyVersion

- RLSubsample: use unified esp_data backend sample_rows() (pandas + polars); drop assumption of raw pandas DataFrame input - _is_empty_labels: use pd.isna() for scalar NaN so legit float labels survive - FillLabelsFromAnswer: replace df.apply with explicit single-pass loop to avoid double-counting during dtype inference; note OOM risk on collect() - build_dataloaders: guard worker_timeout when num_workers == 0 - ebird_taxonomy.load(): document cached mapping as read-only - birdset_train_splits: add missing XCM source - tests: rename misnamed taxonomy test; add RLSubsample + FillLabelsFromAnswer unit tests exercising real esp_data backends

Take the base branch (PR #186) version of transforms.py wholesale — PR #187 carried only a stale snapshot, all divergence is PR #186's RLSubsample backend, _is_empty_labels, and FillLabelsFromAnswer improvements.

GaganNarula · 2026-06-01T11:42:57Z

+            return data, {"subset": self.subset, "original_size": 0, "split_size": 0}
+
+        rng = np.random.default_rng(self.random_state)
+        is_dataframe = isinstance(data, pd.DataFrame)


is there a reason for data to be a DataFrame ? Backwards compat ?

partially. apply_transformations adapts to whatever backend the dataset uses, either polars or pandas

I removed it, it should always be a databackend

GaganNarula

i'm wondering if we actually need audioprotopnet in avex ? if this is just for running an eval then maybe just keep this in a branch ?

* eval: embedding cache metadata + memory fixes - Add _write_embedding_metadata helper; store aggregation in HDF5 attrs - HDF5EmbeddingDataset exposes embedding_aggregation attribute - _embedding_cache_matches validates aggregation before reuse - embedding_manager: use uncompressed size (shape*dtype) for memory routing - __getstate__ strips _window_cache to avoid N_workers x cache duplication - Window loads lazily on first __getitem__ (not prefilled in __init__) - _load_window scales capacity by num_workers to stay within budget * eval: run_evaluate fixes - Use _embedding_cache_matches for cache validation (aggregation-aware) - Fix operator precedence bugs in need_recompute_embeddings_*_clustering - Configurable offline probe embedding aggregation via probe_storage_aggregation - Hoist memory_limit_bytes before probing + retrieval/clustering branches - Graceful per-file fallback when embeddings missing (avoids base_model=None crash) - Release probe HDF5 caches before clustering to free ~46 GB - Reuse probe embeddings for retrieval/clustering when aggregation matches - Pass worker_timeout=120 to build_dataloaders * eval: finetune probe_num_workers + multi-label target fix - Wire probe_num_workers (default 0) into all embedding DataLoaders - Gate pin_memory on probe_workers > 0 (pointless with no workers) - Fix multi-label target shape: one_hot 1D/squeezed-2D before BCEWithLogitsLoss * eval: add probe_storage_aggregation and probe_num_workers to config * fix: skip test_embedding_metadata when esp_data not installed * fix: address Copilot review comments - transforms: add DataFrame path using iloc/unique for pandas compatibility; backend API path retained for esp_data Polars/other backends - pre-commit: narrow large-file and codespell excludes to ebird_taxonomy_v*.json; remove ruff/ruff-format avex/data/ exclude (ruff only targets Python anyway) - ebird_taxonomy: remove unused logger; fix FileNotFoundError docstring indent - birdset_train_splits: add missing Returns/Raises/Yields docstring sections - dataset: document worker_timeout param in build_dataloaders docstring - test_ebird_taxonomy: annotate version param as EbirdTaxonomyVersion * fix: address Copilot review comments - embedding_utils: mark cache incomplete when batches were skipped; skipped rows leave unwritten gaps in pre-allocated HDF5 datasets - run_evaluate: _reuse_probe_embeddings_for_eval rejects aggregation=none (3D embeddings can't be fed into retrieval/clustering) - run_evaluate: replace path.exists() with _embedding_cache_matches for clustering cache; use pooled aggregation (mean/max) for the check * fix: address review comments on eval pipeline - run_evaluate: fix test clustering/retrieval cache never being reused — validate the clustering cache against the pooled clustering aggregation, not the probe storage aggregation (which is 'none' by default) - run_evaluate: introduce shared _POOLED_AGGREGATIONS set so reuse and clustering-recompute agree on cls_token; collapse duplicated force-to-mean blocks to use _clustering_aggregation - run_evaluate: name the eval worker timeout constant (_EVAL_WORKER_TIMEOUT_SECONDS) - embedding_manager: use module-level h5py/numpy instead of local aliased imports - embedding_utils: comment sliding-window sequential-access assumption and the skipped-batch index advance semantics - tests: add _embedding_cache_matches coverage (mismatch, incomplete, legacy attr)

Keep PR #186 backend-only TrainValSplit/RLSubsample/FillLabelsFromAnswer logic; adopt main's non-optional random_state for RLSubsample (#180).

nkundiushuti · 2026-06-30T08:11:17Z

i'm wondering if we actually need audioprotopnet in avex ? if this is just for running an eval then maybe just keep this in a branch ?

@GaganNarula I decided not to add audioprotopnet yet, it's pending some checks and I think we need to separate the prototypical probe head from the audioprotopnet itself. plus it would be good to add it when we have audioprotopnet v2. so if you can approve this I can only merge 1 and 2 and leave 3 for the future

GaganNarula

👍

nkundiushuti requested a review from david-rx as a code owner May 29, 2026 14:52

This was referenced May 29, 2026

[2/3] audioprotopnet: eval pipeline fixes #187

Merged

[1/3] audioprotopnet: data layer updates #188

Closed

[3/3] audioprotopnet: ConvNeXt + AudioProtoPNet + AudioProtoPNetSED models #184

Open

nkundiushuti requested a review from Copilot May 29, 2026 15:59

Copilot started reviewing on behalf of nkundiushuti May 29, 2026 15:59 View session

Copilot AI reviewed May 29, 2026

View reviewed changes

nkundiushuti added 2 commits May 29, 2026 18:15

nkundiushuti requested a review from GaganNarula May 29, 2026 17:21

GaganNarula reviewed Jun 1, 2026

View reviewed changes

assume a databackend from esp_data, add more backend tests

d9c1f69

GaganNarula reviewed Jun 4, 2026

View reviewed changes

nkundiushuti added 2 commits June 30, 2026 10:05

merge: resolve transforms.py conflict with main

237258e

Keep PR #186 backend-only TrainValSplit/RLSubsample/FillLabelsFromAnswer logic; adopt main's non-optional random_state for RLSubsample (#180).

GaganNarula approved these changes Jul 1, 2026

View reviewed changes

nkundiushuti merged commit 7cf81b1 into main Jul 1, 2026
8 checks passed

Uh oh!

Conversation

nkundiushuti commented May 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Dependencies

Test plan

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

GaganNarula Jun 1, 2026

Choose a reason for hiding this comment

Uh oh!

nkundiushuti Jun 2, 2026

Choose a reason for hiding this comment

Uh oh!

nkundiushuti Jun 2, 2026

Choose a reason for hiding this comment

Uh oh!

GaganNarula left a comment

Choose a reason for hiding this comment

Uh oh!

nkundiushuti commented Jun 30, 2026

Uh oh!

GaganNarula left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

nkundiushuti commented May 29, 2026 •

edited

Loading