avex 0.5.0a1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (103) hide show
  1. avex-0.5.0a1/PKG-INFO +259 -0
  2. avex-0.5.0a1/README.md +203 -0
  3. avex-0.5.0a1/avex/__init__.py +51 -0
  4. avex-0.5.0a1/avex/api/__init__.py +1 -0
  5. avex-0.5.0a1/avex/api/configs/__init__.py +0 -0
  6. avex-0.5.0a1/avex/api/configs/official_models/BEATs_iter3_plus_AS2M.yml +17 -0
  7. avex-0.5.0a1/avex/api/configs/official_models/BEATs_iter3_plus_AS2M_finetuned_on_AS2M_cpt2.yml +17 -0
  8. avex-0.5.0a1/avex/api/configs/official_models/__init__.py +1 -0
  9. avex-0.5.0a1/avex/api/configs/official_models/esp_aves2_eat_all.yml +17 -0
  10. avex-0.5.0a1/avex/api/configs/official_models/esp_aves2_eat_bio.yml +17 -0
  11. avex-0.5.0a1/avex/api/configs/official_models/esp_aves2_effnetb0_all.yml +22 -0
  12. avex-0.5.0a1/avex/api/configs/official_models/esp_aves2_effnetb0_audioset.yml +21 -0
  13. avex-0.5.0a1/avex/api/configs/official_models/esp_aves2_effnetb0_bio.yml +22 -0
  14. avex-0.5.0a1/avex/api/configs/official_models/esp_aves2_naturelm_audio_v1_beats.yml +16 -0
  15. avex-0.5.0a1/avex/api/configs/official_models/esp_aves2_sl_beats_all.yml +19 -0
  16. avex-0.5.0a1/avex/api/configs/official_models/esp_aves2_sl_beats_bio.yml +17 -0
  17. avex-0.5.0a1/avex/api/configs/official_models/esp_aves2_sl_eat_all_ssl_all.yml +18 -0
  18. avex-0.5.0a1/avex/api/configs/official_models/esp_aves2_sl_eat_bio_ssl_all.yml +18 -0
  19. avex-0.5.0a1/avex/api/list_models.py +73 -0
  20. avex-0.5.0a1/avex/cli.py +79 -0
  21. avex-0.5.0a1/avex/configs.py +1349 -0
  22. avex-0.5.0a1/avex/data/__init__.py +27 -0
  23. avex-0.5.0a1/avex/data/animalspeak_column_patch.py +124 -0
  24. avex-0.5.0a1/avex/data/audio_utils.py +341 -0
  25. avex-0.5.0a1/avex/data/augmentations.py +541 -0
  26. avex-0.5.0a1/avex/data/cloudpathlib_retry_patch.py +104 -0
  27. avex-0.5.0a1/avex/data/configs.py +232 -0
  28. avex-0.5.0a1/avex/data/data_utils.py +116 -0
  29. avex-0.5.0a1/avex/data/dataset.py +623 -0
  30. avex-0.5.0a1/avex/data/text_label_from_features.py +141 -0
  31. avex-0.5.0a1/avex/data/transforms.py +284 -0
  32. avex-0.5.0a1/avex/evaluation/__init__.py +0 -0
  33. avex-0.5.0a1/avex/evaluation/clustering.py +221 -0
  34. avex-0.5.0a1/avex/evaluation/embedding_manager.py +404 -0
  35. avex-0.5.0a1/avex/evaluation/embedding_utils.py +1572 -0
  36. avex-0.5.0a1/avex/evaluation/finetune.py +799 -0
  37. avex-0.5.0a1/avex/evaluation/retrieval.py +662 -0
  38. avex-0.5.0a1/avex/io/__init__.py +49 -0
  39. avex-0.5.0a1/avex/io/file_utils.py +50 -0
  40. avex-0.5.0a1/avex/io/filesystem.py +89 -0
  41. avex-0.5.0a1/avex/io/paths.py +178 -0
  42. avex-0.5.0a1/avex/metrics/__init__.py +0 -0
  43. avex-0.5.0a1/avex/metrics/beans_metrics.py +324 -0
  44. avex-0.5.0a1/avex/metrics/metric_factory.py +48 -0
  45. avex-0.5.0a1/avex/metrics/sklearn_metrics.py +393 -0
  46. avex-0.5.0a1/avex/models/__init__.py +0 -0
  47. avex-0.5.0a1/avex/models/atst_frame/__init__.py +31 -0
  48. avex-0.5.0a1/avex/models/atst_frame/atst_encoder.py +162 -0
  49. avex-0.5.0a1/avex/models/atst_frame/atst_frame.py +1901 -0
  50. avex-0.5.0a1/avex/models/aves_model.py +280 -0
  51. avex-0.5.0a1/avex/models/base_model.py +387 -0
  52. avex-0.5.0a1/avex/models/beats/__init__.py +30 -0
  53. avex-0.5.0a1/avex/models/beats/backbone.py +906 -0
  54. avex-0.5.0a1/avex/models/beats/beats.py +267 -0
  55. avex-0.5.0a1/avex/models/beats/modules.py +343 -0
  56. avex-0.5.0a1/avex/models/beats_model.py +406 -0
  57. avex-0.5.0a1/avex/models/birdmae.py +267 -0
  58. avex-0.5.0a1/avex/models/birdnet.py +560 -0
  59. avex-0.5.0a1/avex/models/clip.py +189 -0
  60. avex-0.5.0a1/avex/models/eat/__init__.py +4 -0
  61. avex-0.5.0a1/avex/models/eat/audio_processor.py +143 -0
  62. avex-0.5.0a1/avex/models/eat_hf.py +436 -0
  63. avex-0.5.0a1/avex/models/efficientnet.py +342 -0
  64. avex-0.5.0a1/avex/models/get_model.py +182 -0
  65. avex-0.5.0a1/avex/models/perch.py +332 -0
  66. avex-0.5.0a1/avex/models/probes/__init__.py +28 -0
  67. avex-0.5.0a1/avex/models/probes/attention_probe.py +141 -0
  68. avex-0.5.0a1/avex/models/probes/base_probes.py +414 -0
  69. avex-0.5.0a1/avex/models/probes/embedding_projectors.py +565 -0
  70. avex-0.5.0a1/avex/models/probes/linear_probe.py +84 -0
  71. avex-0.5.0a1/avex/models/probes/lstm_probe.py +126 -0
  72. avex-0.5.0a1/avex/models/probes/mlp_probe.py +110 -0
  73. avex-0.5.0a1/avex/models/probes/transformer_probe.py +131 -0
  74. avex-0.5.0a1/avex/models/probes/utils/__init__.py +22 -0
  75. avex-0.5.0a1/avex/models/probes/utils/factory.py +186 -0
  76. avex-0.5.0a1/avex/models/probes/utils/registry.py +93 -0
  77. avex-0.5.0a1/avex/models/resnet.py +157 -0
  78. avex-0.5.0a1/avex/models/surfperch.py +334 -0
  79. avex-0.5.0a1/avex/models/utils/__init__.py +39 -0
  80. avex-0.5.0a1/avex/models/utils/factory.py +165 -0
  81. avex-0.5.0a1/avex/models/utils/load.py +559 -0
  82. avex-0.5.0a1/avex/models/utils/registry.py +694 -0
  83. avex-0.5.0a1/avex/preprocessing/__init__.py +0 -0
  84. avex-0.5.0a1/avex/preprocessing/activity_detector.py +0 -0
  85. avex-0.5.0a1/avex/preprocessing/preprocessor_base.py +0 -0
  86. avex-0.5.0a1/avex/run_evaluate.py +1084 -0
  87. avex-0.5.0a1/avex/run_train.py +181 -0
  88. avex-0.5.0a1/avex/training/__init__.py +0 -0
  89. avex-0.5.0a1/avex/training/checkpoint_manager.py +263 -0
  90. avex-0.5.0a1/avex/training/clustering_evaluator.py +251 -0
  91. avex-0.5.0a1/avex/training/distributed.py +328 -0
  92. avex-0.5.0a1/avex/training/losses.py +358 -0
  93. avex-0.5.0a1/avex/training/metrics_tracker.py +250 -0
  94. avex-0.5.0a1/avex/training/optimisers.py +107 -0
  95. avex-0.5.0a1/avex/training/train.py +751 -0
  96. avex-0.5.0a1/avex/training/trainer_factory.py +180 -0
  97. avex-0.5.0a1/avex/training/training_strategies.py +301 -0
  98. avex-0.5.0a1/avex/training/training_utils.py +74 -0
  99. avex-0.5.0a1/avex/utils/__init__.py +15 -0
  100. avex-0.5.0a1/avex/utils/experiment_logger.py +223 -0
  101. avex-0.5.0a1/avex/utils/experiment_tracking.py +755 -0
  102. avex-0.5.0a1/avex/utils/utils.py +309 -0
  103. avex-0.5.0a1/pyproject.toml +146 -0
avex-0.5.0a1/PKG-INFO ADDED
@@ -0,0 +1,259 @@
1
+ Metadata-Version: 2.3
2
+ Name: avex
3
+ Version: 0.5.0a1
4
+ Summary: A comprehensive Python-based system for training, evaluating, and analyzing audio representation learning models with support for both supervised and self-supervised learning paradigms
5
+ Keywords: audio,machine-learning,avex,self-supervised,bioacoustics
6
+ Author: Earth Species
7
+ Author-email: Earth Species <contact@earthspecies.org>
8
+ License: MIT
9
+ Classifier: Development Status :: 4 - Beta
10
+ Classifier: Intended Audience :: Science/Research
11
+ Classifier: License :: OSI Approved :: MIT License
12
+ Classifier: Operating System :: OS Independent
13
+ Classifier: Programming Language :: Python :: 3
14
+ Classifier: Programming Language :: Python :: 3.10
15
+ Classifier: Programming Language :: Python :: 3.11
16
+ Classifier: Programming Language :: Python :: 3.12
17
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
18
+ Classifier: Topic :: Multimedia :: Sound/Audio :: Analysis
19
+ Requires-Dist: librosa>=0.11.0
20
+ Requires-Dist: numpy>=1.26.4
21
+ Requires-Dist: omegaconf>=2.3.0
22
+ Requires-Dist: pydantic>=2.11.3
23
+ Requires-Dist: pyyaml>=6.0.2
24
+ Requires-Dist: soundfile>=0.13.1
25
+ Requires-Dist: birdnetlib>=0.18.0
26
+ Requires-Dist: timm>=1.0.15
27
+ Requires-Dist: torch>=2.5.0
28
+ Requires-Dist: torchaudio>=2.2.0
29
+ Requires-Dist: torchvision>=0.17.0
30
+ Requires-Dist: tqdm>=4.67.1
31
+ Requires-Dist: huggingface-hub>=0.36.0
32
+ Requires-Dist: transformers>=4.0.0
33
+ Requires-Dist: h5py>=3.13.0
34
+ Requires-Dist: tensorflow>=2.17.0
35
+ Requires-Dist: tensorflow-hub>=0.16.1
36
+ Requires-Dist: einops>=0.8.1
37
+ Requires-Dist: pydantic-settings>=2.9.1
38
+ Requires-Dist: fsspec>=2024.2.0
39
+ Requires-Dist: gcsfs>=2024.2.0
40
+ Requires-Dist: s3fs>=2024.2.0
41
+ Requires-Dist: pandas>=2.2.3
42
+ Requires-Dist: pytorch-lightning>=2.5.1.post0 ; extra == 'dev'
43
+ Requires-Dist: mlflow>=2.22.0 ; extra == 'dev'
44
+ Requires-Dist: wandb>=0.21.0 ; extra == 'dev'
45
+ Requires-Dist: esp-sweep>=0.1.0 ; extra == 'dev'
46
+ Requires-Dist: esp-data>=1.5.0 ; extra == 'dev'
47
+ Requires-Dist: gradio>=4.0.0 ; extra == 'dev'
48
+ Requires-Dist: gradio-leaderboard>=0.0.13 ; extra == 'dev'
49
+ Requires-Python: >=3.10, <3.13
50
+ Project-URL: Homepage, https://github.com/earthspecies/avex
51
+ Project-URL: Repository, https://github.com/earthspecies/avex
52
+ Project-URL: Documentation, https://github.com/earthspecies/avex#readme
53
+ Project-URL: Bug Tracker, https://github.com/earthspecies/avex/issues
54
+ Provides-Extra: dev
55
+ Description-Content-Type: text/markdown
56
+
57
+ # avex - Animal Vocalization Encoder Library
58
+
59
+ ![CI status](https://github.com/earthspecies/avex/actions/workflows/pythonapp.yml/badge.svg?branch=main)
60
+ ![Pre-commit status](https://github.com/earthspecies/avex/actions/workflows/pre-commit.yml/badge.svg?branch=main)
61
+
62
+ An API for model loading and inference, and a Python-based system for training and evaluating bioacoustics representation learning models.
63
+
64
+ ## Description
65
+
66
+ The Animal Vocalization Encoder library avex provides a unified interface for working with pre-trained bioacoustics representation learning models, with support for:
67
+
68
+ - **Model Loading**: Load pre-trained models with checkpoints and class mappings
69
+ - **Embedding Extraction**: Extract features from audio for downstream tasks
70
+ - **Probe System**: Flexible probe heads (linear, MLP, LSTM, attention, transformer) for transfer learning
71
+ - **Training & Evaluation**: Scripts for supervised learning experiments
72
+ - **Plugin Architecture**: Register and use custom models seamlessly
73
+
74
+ ## Installation
75
+
76
+ ### Prerequisites
77
+
78
+ - Python 3.10, 3.11, or 3.12
79
+
80
+ ### Install with pip
81
+
82
+ ```bash
83
+ pip install avex
84
+ ```
85
+
86
+ ### Install with uv
87
+
88
+ ```bash
89
+ uv add avex
90
+ ```
91
+
92
+ For development installation with training/evaluation tools, see the [Contributing guide](CONTRIBUTING.md).
93
+
94
+ ## Quick Start
95
+
96
+ ```python
97
+ import torch
98
+ import librosa
99
+ from avex import load_model, list_models
100
+
101
+ # List available models
102
+ print(list_models().keys())
103
+
104
+ # Load a pre-trained model
105
+ model = load_model("esp_aves2_sl_beats_all", device="cpu")
106
+
107
+ # Load and preprocess audio (BEATs expects 16kHz)
108
+ audio, sr = librosa.load("your_audio.wav", sr=16000)
109
+ audio_tensor = torch.tensor(audio).unsqueeze(0) # Shape: (1, num_samples)
110
+
111
+ # Run inference
112
+ with torch.no_grad():
113
+ logits = model(audio_tensor)
114
+ predicted_class = logits.argmax(dim=-1).item()
115
+
116
+ # Get human-readable label
117
+ if model.label_mapping:
118
+ label = model.label_mapping.get(str(predicted_class), predicted_class)
119
+ print(f"Predicted: {label}")
120
+ ```
121
+
122
+ ### Embedding Extraction
123
+
124
+ ```python
125
+ # Load for embedding extraction (no classifier head)
126
+ model = load_model("esp_aves2_sl_beats_all", return_features_only=True, device="cpu")
127
+
128
+ with torch.no_grad():
129
+ embeddings = model(audio_tensor)
130
+ # Shape: (batch, time_steps, 768) for BEATs
131
+
132
+ # Pool to get fixed-size embedding
133
+ embedding = embeddings.mean(dim=1) # Shape: (batch, 768)
134
+ ```
135
+
136
+ ### Transfer Learning with Probes
137
+
138
+ ```python
139
+ from avex.models.probes import build_probe_from_config
140
+ from avex.configs import ProbeConfig
141
+
142
+ # Load backbone for feature extraction
143
+ base = load_model("esp_aves2_sl_beats_all", return_features_only=True, device="cpu")
144
+
145
+ # Define a probe head for your task
146
+ probe_config = ProbeConfig(
147
+ probe_type="linear",
148
+ target_layers=["last_layer"],
149
+ aggregation="mean",
150
+ freeze_backbone=True,
151
+ online_training=True,
152
+ )
153
+
154
+ probe = build_probe_from_config(
155
+ probe_config=probe_config,
156
+ base_model=base,
157
+ num_classes=10, # Your number of classes
158
+ device="cpu",
159
+ )
160
+ ```
161
+
162
+ ## Documentation
163
+
164
+ **Full documentation**: [docs/index.md](docs/index.md)
165
+
166
+ ### Core Documentation
167
+
168
+ - **[API Reference](docs/api_reference.md)** - Complete API documentation for model loading, registry, and management functions
169
+ - **[Architecture](docs/api_architecture.md)** - Framework architecture, core components, and plugin system
170
+ - **[Supported Models](docs/supported_models.md)** - List of supported models and their configurations
171
+ - **[Configuration](docs/configuration.md)** - ModelSpec parameters, audio requirements, and configuration options
172
+
173
+ ### Usage Guides
174
+
175
+ - **[Training and Evaluation](docs/training_evaluation.md)** - Guide to training and evaluating models
176
+ - **[Embedding Extraction](docs/embedding_extraction.md)** - Working with feature representations and embeddings
177
+ - **[Examples](docs/examples.md)** - Comprehensive examples and use cases
178
+
179
+ ### Advanced Topics
180
+
181
+ - **[Probe System](docs/probe_system.md)** - Understanding and using probes for transfer learning
182
+ - **[API Probes](docs/api_probes.md)** - API reference for probe-related functionality
183
+ - **[Custom Model Registration](docs/custom_model_registration.md)** - Guide on registering custom model classes and loading pre-trained models
184
+
185
+ **Examples**: See the [`examples/`](examples/) directory:
186
+
187
+ - `00_quick_start.py` - Basic model loading
188
+ - `01_basic_model_loading.py` - Loading models with different configurations
189
+ - `02_checkpoint_loading.py` - Working with checkpoints
190
+ - `03_custom_model_registration.py` - Custom model registration
191
+ - `04_training_and_evaluation.py` - Training and evaluation examples
192
+ - `05_embedding_extraction.py` - Feature extraction
193
+ - `06_classifier_head_loading.py` - Classifier head behavior
194
+
195
+ ## Supported Models
196
+
197
+ The framework supports the following audio representation learning models:
198
+
199
+ - **EfficientNet** - EfficientNet-based models for audio classification
200
+ - **BEATs** - BEATs transformer models for audio representation learning
201
+ - **EAT** - Efficient Audio Transformer models
202
+ - **AVES** - AVES model for bioacoustics
203
+ - **BirdMAE** - BirdMAE masked autoencoder for bioacoustic representation learning
204
+ - **ATST** - Audio Spectrogram Transformer
205
+ - **ResNet** - ResNet models (ResNet18, ResNet50, ResNet152)
206
+ - **CLIP** - Contrastive Language-Audio Pretraining models
207
+ - **BirdNet** - BirdNet models for bioacoustic classification
208
+ - **Perch** - Perch models for bioacoustics
209
+ - **SurfPerch** - SurfPerch models
210
+
211
+ See [Supported Models](docs/supported_models.md) for detailed information and configuration examples.
212
+
213
+ ## Supported Probes
214
+
215
+ The framework provides flexible probe heads for transfer learning:
216
+
217
+ - **Linear** - Simple linear classifier (fastest, most memory-efficient)
218
+ - **MLP** - Multi-layer perceptron with configurable hidden layers
219
+ - **LSTM** - Long Short-Term Memory network for sequence modeling
220
+ - **Attention** - Self-attention mechanism for sequence modeling
221
+ - **Transformer** - Full transformer encoder architecture
222
+
223
+ Probes can be trained:
224
+ - **Online**: End-to-end with the backbone (raw audio input)
225
+ - **Offline**: On pre-computed embeddings
226
+
227
+ See [Probe System](docs/probe_system.md) and [API Probes](docs/api_probes.md) for detailed documentation.
228
+
229
+ ## Citing
230
+
231
+ If you use this framework in your research, please cite:
232
+
233
+ ```bibtex
234
+ @article{miron2025matters,
235
+ title={What Matters for Bioacoustic Encoding},
236
+ author={Miron, Marius and Robinson, David and Alizadeh, Milad and Gilsenan-McMahon, Ellen and Narula, Gagan and Pietquin, Olivier and Geist, Matthieu and Chemla, Emmanuel and Cusimano, Maddie and Effenberger, Felix and others},
237
+ journal={arXiv preprint arXiv:2508.11845},
238
+ year={2025}
239
+ }
240
+ ```
241
+
242
+ ## Contributing
243
+
244
+ We welcome contributions! Please see [CONTRIBUTING.md](CONTRIBUTING.md) for:
245
+
246
+ - Development setup
247
+ - Running tests
248
+ - Code style guidelines
249
+ - Adding new functionality
250
+ - Pull request process
251
+
252
+ ## License
253
+
254
+ This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
255
+
256
+ ## Acknowledgments
257
+
258
+ - Built on top of PyTorch
259
+ - Integrates with various pre-trained audio models
avex-0.5.0a1/README.md ADDED
@@ -0,0 +1,203 @@
1
+ # avex - Animal Vocalization Encoder Library
2
+
3
+ ![CI status](https://github.com/earthspecies/avex/actions/workflows/pythonapp.yml/badge.svg?branch=main)
4
+ ![Pre-commit status](https://github.com/earthspecies/avex/actions/workflows/pre-commit.yml/badge.svg?branch=main)
5
+
6
+ An API for model loading and inference, and a Python-based system for training and evaluating bioacoustics representation learning models.
7
+
8
+ ## Description
9
+
10
+ The Animal Vocalization Encoder library avex provides a unified interface for working with pre-trained bioacoustics representation learning models, with support for:
11
+
12
+ - **Model Loading**: Load pre-trained models with checkpoints and class mappings
13
+ - **Embedding Extraction**: Extract features from audio for downstream tasks
14
+ - **Probe System**: Flexible probe heads (linear, MLP, LSTM, attention, transformer) for transfer learning
15
+ - **Training & Evaluation**: Scripts for supervised learning experiments
16
+ - **Plugin Architecture**: Register and use custom models seamlessly
17
+
18
+ ## Installation
19
+
20
+ ### Prerequisites
21
+
22
+ - Python 3.10, 3.11, or 3.12
23
+
24
+ ### Install with pip
25
+
26
+ ```bash
27
+ pip install avex
28
+ ```
29
+
30
+ ### Install with uv
31
+
32
+ ```bash
33
+ uv add avex
34
+ ```
35
+
36
+ For development installation with training/evaluation tools, see the [Contributing guide](CONTRIBUTING.md).
37
+
38
+ ## Quick Start
39
+
40
+ ```python
41
+ import torch
42
+ import librosa
43
+ from avex import load_model, list_models
44
+
45
+ # List available models
46
+ print(list_models().keys())
47
+
48
+ # Load a pre-trained model
49
+ model = load_model("esp_aves2_sl_beats_all", device="cpu")
50
+
51
+ # Load and preprocess audio (BEATs expects 16kHz)
52
+ audio, sr = librosa.load("your_audio.wav", sr=16000)
53
+ audio_tensor = torch.tensor(audio).unsqueeze(0) # Shape: (1, num_samples)
54
+
55
+ # Run inference
56
+ with torch.no_grad():
57
+ logits = model(audio_tensor)
58
+ predicted_class = logits.argmax(dim=-1).item()
59
+
60
+ # Get human-readable label
61
+ if model.label_mapping:
62
+ label = model.label_mapping.get(str(predicted_class), predicted_class)
63
+ print(f"Predicted: {label}")
64
+ ```
65
+
66
+ ### Embedding Extraction
67
+
68
+ ```python
69
+ # Load for embedding extraction (no classifier head)
70
+ model = load_model("esp_aves2_sl_beats_all", return_features_only=True, device="cpu")
71
+
72
+ with torch.no_grad():
73
+ embeddings = model(audio_tensor)
74
+ # Shape: (batch, time_steps, 768) for BEATs
75
+
76
+ # Pool to get fixed-size embedding
77
+ embedding = embeddings.mean(dim=1) # Shape: (batch, 768)
78
+ ```
79
+
80
+ ### Transfer Learning with Probes
81
+
82
+ ```python
83
+ from avex.models.probes import build_probe_from_config
84
+ from avex.configs import ProbeConfig
85
+
86
+ # Load backbone for feature extraction
87
+ base = load_model("esp_aves2_sl_beats_all", return_features_only=True, device="cpu")
88
+
89
+ # Define a probe head for your task
90
+ probe_config = ProbeConfig(
91
+ probe_type="linear",
92
+ target_layers=["last_layer"],
93
+ aggregation="mean",
94
+ freeze_backbone=True,
95
+ online_training=True,
96
+ )
97
+
98
+ probe = build_probe_from_config(
99
+ probe_config=probe_config,
100
+ base_model=base,
101
+ num_classes=10, # Your number of classes
102
+ device="cpu",
103
+ )
104
+ ```
105
+
106
+ ## Documentation
107
+
108
+ **Full documentation**: [docs/index.md](docs/index.md)
109
+
110
+ ### Core Documentation
111
+
112
+ - **[API Reference](docs/api_reference.md)** - Complete API documentation for model loading, registry, and management functions
113
+ - **[Architecture](docs/api_architecture.md)** - Framework architecture, core components, and plugin system
114
+ - **[Supported Models](docs/supported_models.md)** - List of supported models and their configurations
115
+ - **[Configuration](docs/configuration.md)** - ModelSpec parameters, audio requirements, and configuration options
116
+
117
+ ### Usage Guides
118
+
119
+ - **[Training and Evaluation](docs/training_evaluation.md)** - Guide to training and evaluating models
120
+ - **[Embedding Extraction](docs/embedding_extraction.md)** - Working with feature representations and embeddings
121
+ - **[Examples](docs/examples.md)** - Comprehensive examples and use cases
122
+
123
+ ### Advanced Topics
124
+
125
+ - **[Probe System](docs/probe_system.md)** - Understanding and using probes for transfer learning
126
+ - **[API Probes](docs/api_probes.md)** - API reference for probe-related functionality
127
+ - **[Custom Model Registration](docs/custom_model_registration.md)** - Guide on registering custom model classes and loading pre-trained models
128
+
129
+ **Examples**: See the [`examples/`](examples/) directory:
130
+
131
+ - `00_quick_start.py` - Basic model loading
132
+ - `01_basic_model_loading.py` - Loading models with different configurations
133
+ - `02_checkpoint_loading.py` - Working with checkpoints
134
+ - `03_custom_model_registration.py` - Custom model registration
135
+ - `04_training_and_evaluation.py` - Training and evaluation examples
136
+ - `05_embedding_extraction.py` - Feature extraction
137
+ - `06_classifier_head_loading.py` - Classifier head behavior
138
+
139
+ ## Supported Models
140
+
141
+ The framework supports the following audio representation learning models:
142
+
143
+ - **EfficientNet** - EfficientNet-based models for audio classification
144
+ - **BEATs** - BEATs transformer models for audio representation learning
145
+ - **EAT** - Efficient Audio Transformer models
146
+ - **AVES** - AVES model for bioacoustics
147
+ - **BirdMAE** - BirdMAE masked autoencoder for bioacoustic representation learning
148
+ - **ATST** - Audio Spectrogram Transformer
149
+ - **ResNet** - ResNet models (ResNet18, ResNet50, ResNet152)
150
+ - **CLIP** - Contrastive Language-Audio Pretraining models
151
+ - **BirdNet** - BirdNet models for bioacoustic classification
152
+ - **Perch** - Perch models for bioacoustics
153
+ - **SurfPerch** - SurfPerch models
154
+
155
+ See [Supported Models](docs/supported_models.md) for detailed information and configuration examples.
156
+
157
+ ## Supported Probes
158
+
159
+ The framework provides flexible probe heads for transfer learning:
160
+
161
+ - **Linear** - Simple linear classifier (fastest, most memory-efficient)
162
+ - **MLP** - Multi-layer perceptron with configurable hidden layers
163
+ - **LSTM** - Long Short-Term Memory network for sequence modeling
164
+ - **Attention** - Self-attention mechanism for sequence modeling
165
+ - **Transformer** - Full transformer encoder architecture
166
+
167
+ Probes can be trained:
168
+ - **Online**: End-to-end with the backbone (raw audio input)
169
+ - **Offline**: On pre-computed embeddings
170
+
171
+ See [Probe System](docs/probe_system.md) and [API Probes](docs/api_probes.md) for detailed documentation.
172
+
173
+ ## Citing
174
+
175
+ If you use this framework in your research, please cite:
176
+
177
+ ```bibtex
178
+ @article{miron2025matters,
179
+ title={What Matters for Bioacoustic Encoding},
180
+ author={Miron, Marius and Robinson, David and Alizadeh, Milad and Gilsenan-McMahon, Ellen and Narula, Gagan and Pietquin, Olivier and Geist, Matthieu and Chemla, Emmanuel and Cusimano, Maddie and Effenberger, Felix and others},
181
+ journal={arXiv preprint arXiv:2508.11845},
182
+ year={2025}
183
+ }
184
+ ```
185
+
186
+ ## Contributing
187
+
188
+ We welcome contributions! Please see [CONTRIBUTING.md](CONTRIBUTING.md) for:
189
+
190
+ - Development setup
191
+ - Running tests
192
+ - Code style guidelines
193
+ - Adding new functionality
194
+ - Pull request process
195
+
196
+ ## License
197
+
198
+ This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
199
+
200
+ ## Acknowledgments
201
+
202
+ - Built on top of PyTorch
203
+ - Integrates with various pre-trained audio models
@@ -0,0 +1,51 @@
1
+ """
2
+ Representation Learning Framework
3
+
4
+ A comprehensive Python-based system for training, evaluating, and analyzing
5
+ audio representation learning models with support for both supervised and
6
+ self-supervised learning paradigms.
7
+ """
8
+
9
+ from importlib.metadata import version
10
+
11
+ from .models.utils.factory import (
12
+ build_model,
13
+ build_model_from_spec,
14
+ )
15
+ from .models.utils.load import load_label_mapping, load_model
16
+ from .models.utils.registry import (
17
+ describe_model,
18
+ get_checkpoint_path,
19
+ get_model_class,
20
+ get_model_spec,
21
+ list_model_classes,
22
+ list_model_layers,
23
+ list_models,
24
+ register_model,
25
+ register_model_class,
26
+ )
27
+
28
+ __version__ = version("avex")
29
+
30
+
31
+ __all__ = [
32
+ # Model loading
33
+ "load_model",
34
+ # Registry management
35
+ "register_model",
36
+ "get_model_spec",
37
+ "list_models",
38
+ "describe_model",
39
+ "list_model_layers",
40
+ # Model class management
41
+ "register_model_class",
42
+ "get_model_class",
43
+ "list_model_classes",
44
+ # Model factory
45
+ "build_model",
46
+ "build_model_from_spec",
47
+ # Checkpoint management
48
+ "get_checkpoint_path",
49
+ # Label mapping management
50
+ "load_label_mapping",
51
+ ]
@@ -0,0 +1 @@
1
+ """API module for representation learning."""
File without changes
@@ -0,0 +1,17 @@
1
+ ---
2
+ # BEATs iter3+AS2M SSL pretrained backbone
3
+ # Self-supervised learning checkpoint without fine-tuning
4
+ # Weights from HuggingFace: https://huggingface.co/lpepino/beats_ckpts
5
+ checkpoint_path: hf://lpepino/beats_ckpts/BEATs_iter3_plus_AS2M.pt
6
+
7
+ model_spec:
8
+ name: beats
9
+ pretrained: true
10
+ fine_tuned: false
11
+ device: cuda
12
+ audio_config:
13
+ sample_rate: 16000
14
+ representation: raw
15
+ normalize: false
16
+ target_length_seconds: 10
17
+ window_selection: random
@@ -0,0 +1,17 @@
1
+ ---
2
+ # BEATs iter3+AS2M fine-tuned backbone
3
+ # Fine-tuned on AudioSet 2M classification
4
+ # Weights from HuggingFace: https://huggingface.co/THUdyh/Ola_speech_encoders
5
+ checkpoint_path: hf://THUdyh/Ola_speech_encoders/BEATs_iter3_plus_AS2M_finetuned_on_AS2M_cpt2.pt
6
+
7
+ model_spec:
8
+ name: beats
9
+ pretrained: true
10
+ fine_tuned: true
11
+ device: cuda
12
+ audio_config:
13
+ sample_rate: 16000
14
+ representation: raw
15
+ normalize: false
16
+ target_length_seconds: 10
17
+ window_selection: random
@@ -0,0 +1 @@
1
+ """Official model configurations."""
@@ -0,0 +1,17 @@
1
+ ---
2
+ # ESP-AVES2: EAT trained on all (AudioSet + bio) data.
3
+ # Weights on Hugging Face Hub (safetensors). Load via hf:// URIs.
4
+ checkpoint_path: hf://EarthSpeciesProject/esp-aves2-eat-all/esp-aves2-eat-all.safetensors
5
+
6
+ model_spec:
7
+ name: eat_hf
8
+ pretrained: false
9
+ device: cuda
10
+ eat_norm_mean: -5.553
11
+ eat_norm_std: 4.606
12
+ audio_config:
13
+ sample_rate: 16000
14
+ representation: raw
15
+ normalize: false
16
+ target_length_seconds: 10
17
+ window_selection: random
@@ -0,0 +1,17 @@
1
+ ---
2
+ # ESP-AVES2: EAT trained on bio data.
3
+ # Weights on Hugging Face Hub (safetensors). Load via hf:// URIs.
4
+ checkpoint_path: hf://EarthSpeciesProject/esp-aves2-eat-bio/esp-aves2-eat-bio.safetensors
5
+
6
+ model_spec:
7
+ name: eat_hf
8
+ pretrained: false
9
+ device: cuda
10
+ eat_norm_mean: -5.553
11
+ eat_norm_std: 4.606
12
+ audio_config:
13
+ sample_rate: 16000
14
+ representation: raw
15
+ normalize: false
16
+ target_length_seconds: 10
17
+ window_selection: random
@@ -0,0 +1,22 @@
1
+ ---
2
+ # ESP-AVES2: EfficientNet B0 trained on all (AudioSet + bio) data.
3
+ # Weights on Hugging Face Hub (safetensors). Load via hf:// URIs.
4
+ checkpoint_path: hf://EarthSpeciesProject/esp-aves2-effnetb0-all/esp-aves2-effnetb0-all.safetensors
5
+ class_mapping_path: hf://EarthSpeciesProject/esp-aves2-effnetb0-all/label_map.json
6
+
7
+ model_spec:
8
+ name: efficientnet
9
+ pretrained: false
10
+ device: cuda
11
+ efficientnet_variant: b0
12
+ audio_config:
13
+ sample_rate: 16000
14
+ n_fft: 800
15
+ hop_length: 160
16
+ win_length: 800
17
+ window: hann
18
+ n_mels: 128
19
+ representation: mel_spectrogram
20
+ normalize: true
21
+ target_length_seconds: 10
22
+ window_selection: random
@@ -0,0 +1,21 @@
1
+ ---
2
+ # ESP-AVES2: EfficientNet B0 trained on AudioSet.
3
+ # Weights on Hugging Face Hub (safetensors). Load via hf:// URIs.
4
+ checkpoint_path: hf://EarthSpeciesProject/esp-aves2-effnetb0-audioset/esp-aves2-effnetb0-audioset.safetensors
5
+
6
+ model_spec:
7
+ name: efficientnet
8
+ pretrained: false
9
+ device: cuda
10
+ efficientnet_variant: b0
11
+ audio_config:
12
+ sample_rate: 16000
13
+ n_fft: 800
14
+ hop_length: 160
15
+ win_length: 800
16
+ window: hann
17
+ n_mels: 128
18
+ representation: mel_spectrogram
19
+ normalize: true
20
+ target_length_seconds: 10
21
+ window_selection: random
@@ -0,0 +1,22 @@
1
+ ---
2
+ # ESP-AVES2: EfficientNet B0 trained on bio data.
3
+ # Weights on Hugging Face Hub (safetensors). Load via hf:// URIs.
4
+ checkpoint_path: hf://EarthSpeciesProject/esp-aves2-effnetb0-bio/esp-aves2-effnetb0-bio.safetensors
5
+ class_mapping_path: hf://EarthSpeciesProject/esp-aves2-effnetb0-bio/label_map.json
6
+
7
+ model_spec:
8
+ name: efficientnet
9
+ pretrained: false
10
+ device: cuda
11
+ efficientnet_variant: b0
12
+ audio_config:
13
+ sample_rate: 16000
14
+ n_fft: 800
15
+ hop_length: 160
16
+ win_length: 800
17
+ window: hann
18
+ n_mels: 128
19
+ representation: mel_spectrogram
20
+ normalize: true
21
+ target_length_seconds: 10
22
+ window_selection: random
@@ -0,0 +1,16 @@
1
+ ---
2
+ # ESP-AVES2: BEATs with NatureLM (audio v1).
3
+ # Weights on Hugging Face Hub (safetensors). Load via hf:// URIs.
4
+ checkpoint_path: hf://EarthSpeciesProject/esp-aves2-naturelm-audio-v1-beats/esp-aves2-naturelm-audio-v1-beats.safetensors
5
+
6
+ model_spec:
7
+ name: beats
8
+ pretrained: false
9
+ device: cuda
10
+ use_naturelm: true
11
+ audio_config:
12
+ sample_rate: 16000
13
+ representation: raw
14
+ normalize: false
15
+ target_length_seconds: 10
16
+ window_selection: random