avex 0.5.0a1__py3-none-any.whl

avex/__init__.py

@@ -0,0 +1,51 @@
+"""
+Representation Learning Framework
+
+A comprehensive Python-based system for training, evaluating, and analyzing
+audio representation learning models with support for both supervised and
+self-supervised learning paradigms.
+"""
+
+from importlib.metadata import version
+
+from .models.utils.factory import (
+    build_model,
+    build_model_from_spec,
+)
+from .models.utils.load import load_label_mapping, load_model
+from .models.utils.registry import (
+    describe_model,
+    get_checkpoint_path,
+    get_model_class,
+    get_model_spec,
+    list_model_classes,
+    list_model_layers,
+    list_models,
+    register_model,
+    register_model_class,
+)
+
+__version__ = version("avex")
+
+
+__all__ = [
+    # Model loading
+    "load_model",
+    # Registry management
+    "register_model",
+    "get_model_spec",
+    "list_models",
+    "describe_model",
+    "list_model_layers",
+    # Model class management
+    "register_model_class",
+    "get_model_class",
+    "list_model_classes",
+    # Model factory
+    "build_model",
+    "build_model_from_spec",
+    # Checkpoint management
+    "get_checkpoint_path",
+    # Label mapping management
+    "load_label_mapping",
+]

avex/api/__init__.py

@@ -0,0 +1 @@
+"""API module for representation learning."""

avex/api/configs/official_models/BEATs_iter3_plus_AS2M.yml

@@ -0,0 +1,17 @@
+---
+# BEATs iter3+AS2M SSL pretrained backbone
+# Self-supervised learning checkpoint without fine-tuning
+# Weights from HuggingFace: https://huggingface.co/lpepino/beats_ckpts
+checkpoint_path: hf://lpepino/beats_ckpts/BEATs_iter3_plus_AS2M.pt
+
+model_spec:
+  name: beats
+  pretrained: true
+  fine_tuned: false
+  device: cuda
+  audio_config:
+    sample_rate: 16000
+    representation: raw
+    normalize: false
+    target_length_seconds: 10
+    window_selection: random

avex/api/configs/official_models/BEATs_iter3_plus_AS2M_finetuned_on_AS2M_cpt2.yml

@@ -0,0 +1,17 @@
+---
+# BEATs iter3+AS2M fine-tuned backbone
+# Fine-tuned on AudioSet 2M classification
+# Weights from HuggingFace: https://huggingface.co/THUdyh/Ola_speech_encoders
+checkpoint_path: hf://THUdyh/Ola_speech_encoders/BEATs_iter3_plus_AS2M_finetuned_on_AS2M_cpt2.pt
+
+model_spec:
+  name: beats
+  pretrained: true
+  fine_tuned: true
+  device: cuda
+  audio_config:
+    sample_rate: 16000
+    representation: raw
+    normalize: false
+    target_length_seconds: 10
+    window_selection: random

avex/api/configs/official_models/__init__.py

@@ -0,0 +1 @@
+"""Official model configurations."""

avex/api/configs/official_models/esp_aves2_eat_all.yml

@@ -0,0 +1,17 @@
+---
+# ESP-AVES2: EAT trained on all (AudioSet + bio) data.
+# Weights on Hugging Face Hub (safetensors). Load via hf:// URIs.
+checkpoint_path: hf://EarthSpeciesProject/esp-aves2-eat-all/esp-aves2-eat-all.safetensors
+
+model_spec:
+  name: eat_hf
+  pretrained: false
+  device: cuda
+  eat_norm_mean: -5.553
+  eat_norm_std: 4.606
+  audio_config:
+    sample_rate: 16000
+    representation: raw
+    normalize: false
+    target_length_seconds: 10
+    window_selection: random

avex/api/configs/official_models/esp_aves2_eat_bio.yml

@@ -0,0 +1,17 @@
+---
+# ESP-AVES2: EAT trained on bio data.
+# Weights on Hugging Face Hub (safetensors). Load via hf:// URIs.
+checkpoint_path: hf://EarthSpeciesProject/esp-aves2-eat-bio/esp-aves2-eat-bio.safetensors
+
+model_spec:
+  name: eat_hf
+  pretrained: false
+  device: cuda
+  eat_norm_mean: -5.553
+  eat_norm_std: 4.606
+  audio_config:
+    sample_rate: 16000
+    representation: raw
+    normalize: false
+    target_length_seconds: 10
+    window_selection: random

avex/api/configs/official_models/esp_aves2_effnetb0_all.yml

@@ -0,0 +1,22 @@
+---
+# ESP-AVES2: EfficientNet B0 trained on all (AudioSet + bio) data.
+# Weights on Hugging Face Hub (safetensors). Load via hf:// URIs.
+checkpoint_path: hf://EarthSpeciesProject/esp-aves2-effnetb0-all/esp-aves2-effnetb0-all.safetensors
+class_mapping_path: hf://EarthSpeciesProject/esp-aves2-effnetb0-all/label_map.json
+
+model_spec:
+  name: efficientnet
+  pretrained: false
+  device: cuda
+  efficientnet_variant: b0
+  audio_config:
+    sample_rate: 16000
+    n_fft: 800
+    hop_length: 160
+    win_length: 800
+    window: hann
+    n_mels: 128
+    representation: mel_spectrogram
+    normalize: true
+    target_length_seconds: 10
+    window_selection: random

avex/api/configs/official_models/esp_aves2_effnetb0_audioset.yml

@@ -0,0 +1,21 @@
+---
+# ESP-AVES2: EfficientNet B0 trained on AudioSet.
+# Weights on Hugging Face Hub (safetensors). Load via hf:// URIs.
+checkpoint_path: hf://EarthSpeciesProject/esp-aves2-effnetb0-audioset/esp-aves2-effnetb0-audioset.safetensors
+
+model_spec:
+  name: efficientnet
+  pretrained: false
+  device: cuda
+  efficientnet_variant: b0
+  audio_config:
+    sample_rate: 16000
+    n_fft: 800
+    hop_length: 160
+    win_length: 800
+    window: hann
+    n_mels: 128
+    representation: mel_spectrogram
+    normalize: true
+    target_length_seconds: 10
+    window_selection: random

avex/api/configs/official_models/esp_aves2_effnetb0_bio.yml

@@ -0,0 +1,22 @@
+---
+# ESP-AVES2: EfficientNet B0 trained on bio data.
+# Weights on Hugging Face Hub (safetensors). Load via hf:// URIs.
+checkpoint_path: hf://EarthSpeciesProject/esp-aves2-effnetb0-bio/esp-aves2-effnetb0-bio.safetensors
+class_mapping_path: hf://EarthSpeciesProject/esp-aves2-effnetb0-bio/label_map.json
+
+model_spec:
+  name: efficientnet
+  pretrained: false
+  device: cuda
+  efficientnet_variant: b0
+  audio_config:
+    sample_rate: 16000
+    n_fft: 800
+    hop_length: 160
+    win_length: 800
+    window: hann
+    n_mels: 128
+    representation: mel_spectrogram
+    normalize: true
+    target_length_seconds: 10
+    window_selection: random

avex/api/configs/official_models/esp_aves2_naturelm_audio_v1_beats.yml

@@ -0,0 +1,16 @@
+---
+# ESP-AVES2: BEATs with NatureLM (audio v1).
+# Weights on Hugging Face Hub (safetensors). Load via hf:// URIs.
+checkpoint_path: hf://EarthSpeciesProject/esp-aves2-naturelm-audio-v1-beats/esp-aves2-naturelm-audio-v1-beats.safetensors
+
+model_spec:
+  name: beats
+  pretrained: false
+  device: cuda
+  use_naturelm: true
+  audio_config:
+    sample_rate: 16000
+    representation: raw
+    normalize: false
+    target_length_seconds: 10
+    window_selection: random

avex/api/configs/official_models/esp_aves2_sl_beats_all.yml

@@ -0,0 +1,19 @@
+---
+# ESP-AVES2: Supervised/linear-probe BEATs (all) model.
+#
+# Weights are hosted on the Hugging Face Hub as a safetensors file.
+# The loader supports `hf://` URIs via huggingface_hub's fsspec filesystem.
+checkpoint_path: hf://EarthSpeciesProject/esp-aves2-sl-beats-all/esp-aves2-sl-beats-all.safetensors
+class_mapping_path: hf://EarthSpeciesProject/esp-aves2-sl-beats-all/label_map.json
+
+model_spec:
+  name: beats
+  pretrained: false
+  fine_tuned: true
+  device: cuda
+  audio_config:
+    sample_rate: 16000
+    representation: raw
+    normalize: false
+    target_length_seconds: 10
+    window_selection: random

avex/api/configs/official_models/esp_aves2_sl_beats_bio.yml

@@ -0,0 +1,17 @@
+---
+# ESP-AVES2: Supervised/linear-probe BEATs (bio) model.
+# Weights on Hugging Face Hub (safetensors). Load via hf:// URIs.
+checkpoint_path: hf://EarthSpeciesProject/esp-aves2-sl-beats-bio/esp-aves2-sl-beats-bio.safetensors
+class_mapping_path: hf://EarthSpeciesProject/esp-aves2-sl-beats-bio/label_map.json
+
+model_spec:
+  name: beats
+  pretrained: false
+  fine_tuned: true
+  device: cuda
+  audio_config:
+    sample_rate: 16000
+    representation: raw
+    normalize: false
+    target_length_seconds: 10
+    window_selection: random

avex/api/configs/official_models/esp_aves2_sl_eat_all_ssl_all.yml

@@ -0,0 +1,18 @@
+---
+# ESP-AVES2: EAT SSL (all data) + linear head.
+# Weights on Hugging Face Hub (safetensors). Load via hf:// URIs.
+checkpoint_path: hf://EarthSpeciesProject/esp-aves2-sl-eat-all-ssl-all/esp-aves2-sl-eat-all-ssl-all.safetensors
+class_mapping_path: hf://EarthSpeciesProject/esp-aves2-sl-eat-all-ssl-all/label_map.json
+
+model_spec:
+  name: eat_hf
+  pretrained: false
+  device: cuda
+  eat_norm_mean: -5.553
+  eat_norm_std: 4.606
+  audio_config:
+    sample_rate: 16000
+    representation: raw
+    normalize: false
+    target_length_seconds: 10
+    window_selection: random

avex/api/configs/official_models/esp_aves2_sl_eat_bio_ssl_all.yml

@@ -0,0 +1,18 @@
+---
+# ESP-AVES2: EAT SSL (bio) + linear head.
+# Weights on Hugging Face Hub (safetensors). Load via hf:// URIs.
+checkpoint_path: hf://EarthSpeciesProject/esp-aves2-sl-eat-bio-ssl-all/esp-aves2-sl-eat-bio-ssl-all.safetensors
+class_mapping_path: hf://EarthSpeciesProject/esp-aves2-sl-eat-bio-ssl-all/label_map.json
+
+model_spec:
+  name: eat_hf
+  pretrained: false
+  device: cuda
+  eat_norm_mean: -5.553
+  eat_norm_std: 4.606
+  audio_config:
+    sample_rate: 16000
+    representation: raw
+    normalize: false
+    target_length_seconds: 10
+    window_selection: random

avex/api/list_models.py

@@ -0,0 +1,73 @@
+#!/usr/bin/env python3
+"""
+CLI helper to list available models in the registry.
+
+Usage:
+    python -m avex.api.list_models
+    python -m avex.api.list_models --detailed
+"""
+
+import sys
+
+import click
+
+from avex.models.utils.registry import list_models
+
+
+def print_models(detailed: bool = False) -> None:
+    """Print available models in the registry.
+
+    Args:
+        detailed: If True, show detailed model information
+    """
+    models = list_models()
+
+    if not models:
+        print("No models registered in the registry.")
+        return
+
+    print(f"Available models ({len(models)}):")
+    print("=" * 50)
+
+    for name, model_spec in models.items():
+        print(f"\n📦 {name}")
+
+        if detailed:
+            print(f"  Name: {model_spec.name}")
+            print(f"  Pretrained: {model_spec.pretrained}")
+            print(f"  Device: {model_spec.device}")
+
+            if hasattr(model_spec, "audio_config") and model_spec.audio_config:
+                audio_cfg = model_spec.audio_config
+                print("  Audio Config:")
+                print(f"    Sample Rate: {audio_cfg.sample_rate}")
+                print(f"    Representation: {audio_cfg.representation}")
+                print(f"    Target Length: {audio_cfg.target_length_seconds}s")
+                print(f"    Normalize: {audio_cfg.normalize}")
+
+            if hasattr(model_spec, "use_naturelm"):
+                print(f"  NatureLM: {model_spec.use_naturelm}")
+        else:
+            # Simple format
+            pretrained_str = "✓" if model_spec.pretrained else "✗"
+            print(f"  Pretrained: {pretrained_str} | Device: {model_spec.device}")
+
+
+@click.command()
+@click.option(
+    "--detailed",
+    "-d",
+    is_flag=True,
+    help="Show detailed model information",
+)
+def main(detailed: bool) -> None:
+    """List available models in the representation learning registry."""
+    try:
+        print_models(detailed=detailed)
+    except Exception as e:
+        click.echo(f"Error listing models: {e}", err=True)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

avex/cli.py

@@ -0,0 +1,79 @@
+"""
+Command-line interface for representation learning experiments.
+
+This module provides a unified CLI for both training and evaluation tasks.
+"""
+
+import logging
+import os
+from pathlib import Path
+
+import click
+
+# Get log level from environment variable, default to INFO
+log_level = os.environ.get("LOG_LEVEL", "INFO").upper()
+numeric_level = getattr(logging, log_level, logging.INFO)
+
+logging.basicConfig(
+    level=numeric_level,
+    format="%(asctime)s | %(levelname)s | %(name)s: %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+)
+
+logger = logging.getLogger(__name__)
+
+
+@click.group()
+def cli() -> None:
+    """Representation learning CLI for training and evaluation."""
+    pass
+
+
+@cli.command()
+@click.option(
+    "--config",
+    "-c",
+    type=click.Path(exists=True, path_type=Path),
+    default="configs/run_configs/clip_base.yml",
+    help="Path to the training config file",
+)
+@click.option(
+    "--patch",
+    "-p",
+    "patches",
+    type=str,
+    multiple=True,
+    help="Patch config values in format 'key=value'. Can be used multiple times.",
+)
+def train(config: Path, patches: tuple[str, ...]) -> None:
+    """Train an audio representation model."""
+    from avex.run_train import main as train_main
+
+    train_main(config_path=config, patches=patches)
+
+
+@cli.command()
+@click.option(
+    "--config",
+    "-c",
+    type=click.Path(exists=True, path_type=Path),
+    required=True,
+    help="Path to evaluation config file",
+)
+@click.option(
+    "--patch",
+    "-p",
+    "patches",
+    type=str,
+    multiple=True,
+    help="Patch config values in format 'key=value'. Can be used multiple times.",
+)
+def evaluate(config: Path, patches: tuple[str, ...]) -> None:
+    """Run linear-probe / fine-tuning experiments."""
+    from avex.run_evaluate import main as eval_main
+
+    eval_main(config_path=config, patches=patches)
+
+
+if __name__ == "__main__":
+    cli()