bioamla.ml

bioamla.ml

bioamla.ml — machine-learning domain.

Audio Spectrogram Transformer (AST) inference, training, and embeddings, on top of the device / base-model foundations.

PyTorch / torchaudio / transformers ship in the base install but are imported lazily inside functions/methods so this package imports fast. Load / inference failures raise :class:~bioamla.exceptions.ModelError.

Example

from bioamla.ml import ASTInference inference = ASTInference(model_path="bioamla/scp-frogs") result = inference.predict("audio.wav") print(result.predicted_label, result.confidence)

InferenceConfig dataclass

Configuration for optimized batch inference.

ASTModel

Bases: BaseAudioModel

Audio Spectrogram Transformer model wrapper.

Provides a unified interface for AST models from the HuggingFace transformers library.

Example

model = ASTModel() model.load("MIT/ast-finetuned-audioset-10-10-0.4593") results = model.predict("audio.wav")

__init__

__init__(config: ModelConfig | None = None) -> None

Initialize AST model.

load

load(
    model_path: str,
    use_fp16: bool = False,
    use_compile: bool = False,
    **kwargs,
) -> ASTModel

Load an AST model from a path or the HuggingFace Hub.

Parameters:

Name	Type	Description	Default
model_path	str	Path to the model or a HuggingFace model identifier.	required
use_fp16	bool	Use half-precision inference.	False
use_compile	bool	Wrap the model with torch.compile().	False

Returns:

Type	Description
ASTModel	Self, for method chaining.

Raises:

Type	Description
ModelError	If the model cannot be loaded.

predict

predict(
    audio: Union[str, ndarray, Tensor],
    sample_rate: int | None = None,
) -> list[PredictionResult]

Run prediction on audio, returning one result per segment.

Parameters:

Name	Type	Description	Default
audio	Union[str, ndarray, Tensor]	Audio file path, numpy array, or torch tensor.	required
sample_rate	int | None	Sample rate if audio is an array/tensor.	None

Returns:

Type	Description
list[PredictionResult]	List of prediction results.

Raises:

Type	Description
ModelError	If the model is not loaded or inference fails.

extract_embeddings

extract_embeddings(
    audio: Union[str, ndarray, Tensor],
    sample_rate: int | None = None,
    layer: str | None = None,
) -> np.ndarray

Extract embeddings from audio (mean-pooled hidden state).

Parameters:

Name	Type	Description	Default
audio	Union[str, ndarray, Tensor]	Audio file path, numpy array, or torch tensor.	required
sample_rate	int | None	Sample rate if audio is an array/tensor.	None
layer	str | None	Layer to extract from (last_hidden_state, layer_<n>, or default = mean of the last hidden state).	None

Returns:

Type	Description
ndarray	Embedding vectors as a numpy array.

Raises:

Type	Description
ModelError	If the model is not loaded or inference fails.

get_attention_weights

get_attention_weights(
    audio: Union[str, ndarray, Tensor],
    sample_rate: int | None = None,
) -> list[np.ndarray]

Get per-layer attention weight matrices for the audio.

Parameters:

Name	Type	Description	Default
audio	Union[str, ndarray, Tensor]	Audio file path, numpy array, or torch tensor.	required
sample_rate	int | None	Sample rate if audio is an array/tensor.	None

Returns:

Type	Description
list[ndarray]	List of attention weight matrices, one per layer.

Raises:

Type	Description
ModelError	If the model is not loaded or inference fails.

EvaluationResult dataclass

Result of an AST model evaluation over a labelled directory.

to_dict

to_dict() -> dict[str, Any]

Convert to a plain dict.

TrainResult dataclass

Result of an AST model training run.

to_dict

to_dict() -> dict[str, Any]

Convert to a plain dict.

BaseAudioModel

Bases: ABC

Abstract base class for audio classification models.

This class defines the interface that all model backends must implement. It provides common functionality for audio preprocessing, batch processing, and result filtering.

backend abstractmethod property

backend: ModelBackend

Return the model backend type.

num_classes property

num_classes: int

Return the number of classes.

classes property

classes: list[str]

Return list of class labels.

__init__

__init__(config: ModelConfig | None = None)

Initialize the model.

Parameters:

Name	Type	Description	Default
config	ModelConfig | None	Model configuration. Uses defaults if None.	None

load abstractmethod

load(model_path: str, **kwargs: Any) -> BaseAudioModel

Load model from path.

Parameters:

Name	Type	Description	Default
model_path	str	Path to model file or HuggingFace identifier.	required
**kwargs	Any	Additional model-specific arguments.	{}

Returns:

Type	Description
BaseAudioModel	Self for method chaining.

predict abstractmethod

predict(
    audio: Union[str, ndarray, Tensor],
    sample_rate: int | None = None,
) -> list[PredictionResult]

Run prediction on audio.

Parameters:

Name	Type	Description	Default
audio	Union[str, ndarray, Tensor]	Audio file path, numpy array, or torch tensor.	required
sample_rate	int | None	Sample rate if audio is an array/tensor.	None

Returns:

Type	Description
list[PredictionResult]	List of prediction results.

predict_file

predict_file(filepath: str) -> list[PredictionResult]

Run prediction on an audio file.

Parameters:

Name	Type	Description	Default
filepath	str	Path to the audio file.	required

Returns:

Type	Description
list[PredictionResult]	List of prediction results.

extract_embeddings abstractmethod

extract_embeddings(
    audio: Union[str, ndarray, Tensor],
    sample_rate: int | None = None,
    layer: str | None = None,
) -> np.ndarray

Extract embeddings from audio.

Parameters:

Name	Type	Description	Default
audio	Union[str, ndarray, Tensor]	Audio file path, numpy array, or torch tensor.	required
sample_rate	int | None	Sample rate if audio is an array/tensor.	None
layer	str | None	Optional layer name for extraction.	None

Returns:

Type	Description
ndarray	Embedding vectors as numpy array.

predict_batch

predict_batch(
    audio_files: list[str],
    progress_callback: Callable[[int, int], None]
    | None = None,
) -> BatchPredictionResult

Run batch prediction on multiple files.

Parameters:

Name	Type	Description	Default
audio_files	list[str]	List of audio file paths.	required
progress_callback	Callable[[int, int], None] | None	Optional callback(current, total) for progress.	None

Returns:

Type	Description
BatchPredictionResult	Batch prediction results.

filter_predictions

filter_predictions(
    predictions: list[PredictionResult],
    min_confidence: float | None = None,
    labels: list[str] | None = None,
    exclude_labels: list[str] | None = None,
) -> list[PredictionResult]

Filter predictions by confidence and labels.

Parameters:

Name	Type	Description	Default
predictions	list[PredictionResult]	List of predictions to filter.	required
min_confidence	float | None	Minimum confidence threshold.	None
labels	list[str] | None	Only include these labels.	None
exclude_labels	list[str] | None	Exclude these labels.	None

Returns:

Type	Description
list[PredictionResult]	Filtered list of predictions.

save

save(path: str, format: str = 'pt') -> str

Save model to file.

Parameters:

Name	Type	Description	Default
path	str	Output path.	required
format	str	Save format ("pt" for PyTorch, "onnx" for ONNX).	'pt'

Returns:

Type	Description
str	Path to saved model.

to

to(device: Union[str, device]) -> BaseAudioModel

Move model to device.

eval

eval() -> BaseAudioModel

Set model to evaluation mode.

BatchPredictionResult dataclass

Results from batch prediction.

to_dict

to_dict() -> dict[str, Any]

Convert to dictionary.

ModelBackend

Bases: Enum

Supported model backends.

ModelConfig dataclass

Base configuration for all models.

get_device

get_device() -> torch.device

Get the torch device.

PredictionResult dataclass

Result from a single prediction.

to_dict

to_dict() -> dict[str, Any]

Convert to dictionary.

DeviceContext

Context manager for temporarily using a specific device.

Example

with DeviceContext("cuda:1"): # Operations here use cuda:1 pass

Back to original device

BatchEmbeddingResult dataclass

Result from batch embedding extraction.

embedding_dim property

embedding_dim: int

Embedding dimension.

success_rate property

success_rate: float

Percentage of files successfully processed.

EmbeddingConfig dataclass

Configuration for embedding extraction.

EmbeddingExtractor

Unified embedding extractor (AST backend).

Example

extractor = EmbeddingExtractor(model_path="MIT/ast-finetuned-audioset-10-10-0.4593") result = extractor.extract("audio.wav") print(result.embeddings.shape)

__init__

__init__(
    model_path: str | None = None,
    config: EmbeddingConfig | None = None,
)

Initialize the embedding extractor.

Parameters:

Name	Type	Description	Default
model_path	str | None	Path to model or HuggingFace model ID.	None
config	EmbeddingConfig | None	Embedding configuration (uses defaults if None).	None

extract

extract(
    audio: str | ndarray,
    sample_rate: int | None = None,
    layer: str | None = None,
) -> EmbeddingResult

Extract embeddings from a single audio file or array.

Parameters:

Name	Type	Description	Default
audio	str | ndarray	Audio file path or numpy array.	required
sample_rate	int | None	Sample rate (required if audio is an array).	None
layer	str | None	Layer to extract from (uses config default if None).	None

Returns:

Name	Type	Description
An	EmbeddingResult	class:EmbeddingResult.

Raises:

Type	Description
ModelError	If extraction fails.

extract_iter

extract_iter(
    audio_files: list[str],
    progress_callback: Callable[[int, int], None]
    | None = None,
) -> Iterator[EmbeddingResult]

Extract embeddings from multiple files, yielding per-file results.

Files that fail are yielded with an empty embeddings array and an error entry in their metadata (the iterator does not raise).

Parameters:

Name	Type	Description	Default
audio_files	list[str]	Audio file paths.	required
progress_callback	Callable[[int, int], None] | None	Optional (current, total) callback.	None

Yields:

Name	Type	Description
An	EmbeddingResult	class:EmbeddingResult for each file.

extract_batch

extract_batch(
    audio_files: list[str],
    aggregate: str = "mean",
    progress_callback: Callable[[int, int], None]
    | None = None,
) -> BatchEmbeddingResult

Extract embeddings from multiple files into a single stacked array.

Parameters:

Name	Type	Description	Default
audio_files	list[str]	Audio file paths.	required
aggregate	str	Segment aggregation: "mean", "first", or "all".	'mean'
progress_callback	Callable[[int, int], None] | None	Optional (current, total) callback.	None

Returns:

Name	Type	Description
A	BatchEmbeddingResult	class:BatchEmbeddingResult.

fit_reducer

fit_reducer(embeddings: ndarray) -> None

Fit the configured dimensionality reducer on embeddings.

EmbeddingResult dataclass

Result from single-file embedding extraction.

embedding_dim property

embedding_dim: int

Embedding dimension.

num_segments property

num_segments: int

Number of segments (embeddings).

mean_embedding

mean_embedding() -> np.ndarray

Get mean embedding across all segments.

ASTInference

High-level inference engine for AST models.

Handles model loading, feature extraction, and prediction for audio classification.

Raises:

Type	Description
ModelError	If the model cannot be loaded.

__init__

__init__(
    model_path: str,
    sample_rate: int = 16000,
    device: Union[str, device] | None = None,
)

Initialize the inference engine and load the model.

Parameters:

Name	Type	Description	Default
model_path	str	Path to the trained AST model or a HuggingFace identifier.	required
sample_rate	int	Target sample rate for audio preprocessing.	16000
device	Union[str, device] | None	Inference device (auto-detected if None).	None

predict

predict(
    audio_path: str, return_logits: bool = False
) -> ASTPredictionResult

Run inference on a single audio file (treated as one clip).

Parameters:

Name	Type	Description	Default
audio_path	str	Path to the audio file.	required
return_logits	bool	Include raw logits in the result if True.	False

Returns:

Name	Type	Description
An	ASTPredictionResult	class:ASTPredictionResult.

Raises:

Type	Description
ModelError	If loading the audio or running inference fails.

predict_topk

predict_topk(
    audio_path: str,
    *,
    top_k: int = 5,
    min_confidence: float = 0.0,
) -> ASTPredictionResult

Run inference on a single file, returning the top-k labels and scores.

The top-1 label/confidence populate predicted_label/confidence as usual; top_k_labels / top_k_scores hold the ranked top-k predictions whose probability is at least min_confidence.

Parameters:

Name	Type	Description	Default
audio_path	str	Path to the audio file.	required
top_k	int	Number of top predictions to keep.	5
min_confidence	float	Drop predictions below this probability.	0.0

Returns:

Name	Type	Description
An	ASTPredictionResult	class:ASTPredictionResult with top-k fields populated.

Raises:

Type	Description
ModelError	If loading the audio or running inference fails.

predict_segments

predict_segments(
    audio_path: str,
    clip_length: int = 10,
    overlap: int = 0,
    return_logits: bool = False,
) -> list[ASTPredictionResult]

Run inference on an audio file split into segments.

Parameters:

Name	Type	Description	Default
audio_path	str	Path to the audio file.	required
clip_length	int	Duration of each segment in seconds.	10
overlap	int	Overlap between segments in seconds.	0
return_logits	bool	Include raw logits in each result if True.	False

Returns:

Type	Description
list[ASTPredictionResult]	A list of :class:ASTPredictionResult, one per segment.

Raises:

Type	Description
ModelError	If loading the audio or running inference fails.

ASTPredictionResult dataclass

Result from a single AST inference (one clip or segment).

Distinct from :class:bioamla.ml.base.PredictionResult; this is the flat, inference-oriented shape returned by :class:ASTInference and consumed by the models ast predict CLI command.

AugmentationConfig dataclass

Configuration for spectrogram augmentation.

Attributes:

Name	Type	Description
time_mask	bool	Enable time masking (SpecAugment).
frequency_mask	bool	Enable frequency masking (SpecAugment).
time_mask_max_masks	int	Maximum number of time masks.
time_mask_max_length	float	Maximum length of time mask (fraction of total).
frequency_mask_max_masks	int	Maximum number of frequency masks.
frequency_mask_max_length	float	Maximum length of frequency mask (fraction of total).
random_gain	bool	Enable random gain adjustment.
gain_range_db	tuple[float, float]	Gain range in dB (min, max).
add_noise	bool	Enable adding background noise.

BioamlaPreprocessor

Mel-spectrogram preprocessing for AST, with optional SpecAugment.

Generates mel spectrograms from audio files or raw samples using librosa. When augmentation is enabled (training), applies audio-domain gain/noise and spectrogram-domain time/frequency masking. process_samples never augments (inference path).

Example

preprocessor = BioamlaPreprocessor(sample_duration=3.0, sample_rate=16000) spectrogram = preprocessor.process_file("audio.wav") spectrogram.shape (128, 94) # (n_mels, time_frames)

With augmentation for training

from bioamla.ml import AugmentationConfig aug_config = AugmentationConfig(time_mask=True, frequency_mask=True) preprocessor.enable_augmentation(aug_config) augmented_spec = preprocessor.process_file("audio.wav")

augmentation_enabled property

augmentation_enabled: bool

Check if augmentation is enabled.

config property

config: dict[str, Any]

Get preprocessor configuration.

__init__

__init__(
    sample_duration: float = 3.0,
    sample_rate: int = 16000,
    n_mels: int = 128,
    n_fft: int = 2048,
    hop_length: int = 512,
    f_min: float = 0.0,
    f_max: float | None = None,
    height: int | None = None,
    width: int | None = None,
) -> None

Initialize the preprocessor.

Parameters:

Name	Type	Description	Default
sample_duration	float	Duration of audio clips in seconds.	3.0
sample_rate	int	Target sample rate in Hz.	16000
n_mels	int	Number of mel bands.	128
n_fft	int	FFT window size.	2048
hop_length	int	Hop length for STFT.	512
f_min	float	Minimum frequency for mel filterbank.	0.0
f_max	float | None	Maximum frequency (None = Nyquist).	None
height	int | None	Output spectrogram height (None = n_mels).	None
width	int | None	Output spectrogram width (None = computed from duration).	None

enable_augmentation

enable_augmentation(config: AugmentationConfig) -> None

Enable augmentation with the given configuration.

disable_augmentation

disable_augmentation() -> None

Disable all augmentations.

process_file

process_file(
    filepath: str,
    start_time: float | None = None,
    end_time: float | None = None,
) -> np.ndarray

Process an audio file to generate a mel spectrogram.

Applies augmentation when enabled (audio-domain on samples, then spectrogram-domain masking on the mel).

Parameters:

Name	Type	Description	Default
filepath	str	Path to audio file.	required
start_time	float | None	Optional start time in seconds (default: 0).	None
end_time	float | None	Optional end time in seconds (default: start + sample_duration).	None

Returns:

Type	Description
ndarray	Mel spectrogram as 2D numpy array (frequency x time).

process_samples

process_samples(
    samples: ndarray, sample_rate: int
) -> np.ndarray

Process raw audio samples to generate a mel spectrogram.

Augmentation is not applied when processing samples directly.

Parameters:

Name	Type	Description	Default
samples	ndarray	Audio samples as 1D numpy array.	required
sample_rate	int	Sample rate of the audio.	required

Returns:

Type	Description
ndarray	Mel spectrogram as 2D numpy array (frequency x time).

to_tensor

to_tensor(
    spectrogram: ndarray, normalize: bool = True
) -> torch.Tensor

Convert spectrogram to a PyTorch tensor.

Parameters:

Name	Type	Description	Default
spectrogram	ndarray	Input spectrogram as numpy array.	required
normalize	bool	Normalize to [0, 1] range.	True

Returns:

Type	Description
Tensor	Spectrogram as a PyTorch tensor with a leading channel dim.

ast_predict

ast_predict(
    input_values: Tensor,
    model: AutoModelForAudioClassification,
) -> str

Run an AST model on preprocessed features and return the predicted label.

Parameters:

Name	Type	Description	Default
input_values	Tensor	Preprocessed audio features from the feature extractor.	required
model	AutoModelForAudioClassification	The AST model to use for prediction.	required

Returns:

Type	Description
str	The predicted class label.

Raises:

Type	Description
ModelError	If inference fails.

ast_predict_batch

ast_predict_batch(
    input_values: Tensor,
    model: AutoModelForAudioClassification,
) -> list[str]

Run an AST model on a batch of preprocessed features.

Parameters:

Name	Type	Description	Default
input_values	Tensor	Batched preprocessed audio features.	required
model	AutoModelForAudioClassification	The AST model to use for prediction.	required

Returns:

Type	Description
list[str]	Predicted class labels, one per item in the batch.

Raises:

Type	Description
ModelError	If inference fails.

extract_features

extract_features(
    waveform_tensor: Tensor,
    sample_rate: int,
    feature_extractor: Optional[ASTFeatureExtractor] = None,
    device: Optional[device] = None,
) -> torch.Tensor

Extract AST input features from an audio waveform tensor.

Parameters:

Name	Type	Description	Default
waveform_tensor	Tensor	Audio waveform as a tensor.	required
sample_rate	int	Sampling rate of the audio.	required
feature_extractor	Optional[ASTFeatureExtractor]	Optional cached feature extractor.	None
device	Optional[device]	Optional device for the output tensor.	None

Returns:

Type	Description
Tensor	Extracted features (input_values) ready for model input.

get_cached_feature_extractor cached

get_cached_feature_extractor(
    model_path: str | None = None,
) -> ASTFeatureExtractor

Get a cached AST feature extractor instance.

Uses an LRU cache to avoid recreating the feature extractor on every call. If loading from a model path fails (e.g. missing preprocessor_config.json), falls back to the default :class:ASTFeatureExtractor.

Parameters:

Name	Type	Description	Default
model_path	str | None	Optional path to load a feature extractor from a specific model. If None or if loading fails, uses the default extractor.	None

Returns:

Name	Type	Description
ASTFeatureExtractor	ASTFeatureExtractor	Cached feature extractor instance.

load_pretrained_ast_model

load_pretrained_ast_model(
    model_path: str,
    use_fp16: bool = False,
    use_compile: bool = False,
) -> AutoModelForAudioClassification

Load a pre-trained AST model from a path or HuggingFace identifier.

Parameters:

Name	Type	Description	Default
model_path	str	Path to the model directory or HuggingFace model identifier.	required
use_fp16	bool	If True, load the model in half precision (FP16).	False
use_compile	bool	If True, wrap the model with torch.compile() (PyTorch 2.0+).	False

Returns:

Type	Description
AutoModelForAudioClassification	The loaded AST model, ready for inference.

Raises:

Type	Description
ModelError	If the model cannot be loaded.

wav_ast_inference

wav_ast_inference(
    wave_path: str, model_path: str, sample_rate: int
) -> str

Run AST inference on a single audio file and return one prediction.

Loads the file, resamples, loads the model, and returns a single predicted label for the entire file.

Parameters:

Name	Type	Description	Default
wave_path	str	Path to the audio file to classify.	required
model_path	str	Path to the pre-trained AST model.	required
sample_rate	int	Target sampling rate for preprocessing.	required

Returns:

Type	Description
str	The predicted class label.

Raises:

Type	Description
ModelError	If loading or inference fails.

evaluate_directory

evaluate_directory(
    audio_dir: str,
    model_path: str,
    ground_truth_csv: str,
    file_column: str = "file_name",
    label_column: str = "label",
    resample_freq: int = 16000,
    use_fp16: bool = False,
) -> EvaluationResult

Evaluate an AST model over a directory of audio files with ground truth.

Parameters:

Name	Type	Description	Default
audio_dir	str	Directory containing audio files.	required
model_path	str	Path to model or HuggingFace identifier.	required
ground_truth_csv	str	CSV with ground-truth labels.	required
file_column	str	Column name for file names in the CSV.	'file_name'
label_column	str	Column name for labels in the CSV.	'label'
resample_freq	int	Target sample rate.	16000
use_fp16	bool	Use half-precision inference.	False

Returns:

Name	Type	Description
An	EvaluationResult	class:EvaluationResult.

Raises:

Type	Description
NotFoundError	If a required path or matching audio is missing.
InvalidInputError	If the CSV is malformed.
ModelError	On model-load or inference failure.

extract_embeddings_file

extract_embeddings_file(
    filepath: str,
    model_path: str,
    layer: str | None = None,
    sample_rate: int = 16000,
) -> dict[str, Any]

Extract AST embeddings from a single file (CLS-token of the base model).

Mirrors the old ASTService.extract_embeddings behavior: uses AutoModel and the CLS token of the last hidden state.

Parameters:

Name	Type	Description	Default
filepath	str	Path to the audio file.	required
model_path	str	Path to model or HuggingFace identifier.	required
layer	str | None	Unused placeholder kept for API compatibility.	None
sample_rate	int	Target sample rate.	16000

Returns:

Type	Description
dict[str, Any]	A dict with filepath, embeddings (numpy array), shape, and
dict[str, Any]	dtype.

Raises:

Type	Description
NotFoundError	If the file does not exist.
ModelError	On model-load or inference failure.

get_model_info

get_model_info(model_path: str) -> dict[str, Any]

Return lightweight info about an AST model from its config.

Parameters:

Name	Type	Description	Default
model_path	str	Path to model or HuggingFace identifier.	required

Returns:

Type	Description
dict[str, Any]	A dict with path, model_type, num_classes, classes,
dict[str, Any]	has_more_classes, and hidden_size.

Raises:

Type	Description
ModelError	On model-load or inference failure.

predict_file

predict_file(
    filepath: str,
    model_path: str = "bioamla/scp-frogs",
    resample_freq: int = 16000,
) -> ASTPredictionResult

Run AST prediction on a single audio file (whole-file).

Parameters:

Name	Type	Description	Default
filepath	str	Path to the audio file.	required
model_path	str	Path to model or HuggingFace identifier.	'bioamla/scp-frogs'
resample_freq	int	Target sample rate.	16000

Returns:

Name	Type	Description
An	ASTPredictionResult	class:~bioamla.ml.inference.ASTPredictionResult.

Raises:

Type	Description
NotFoundError	If the file does not exist.
ModelError	On inference failure.

AudioDataset

AudioDataset(
    filepaths: list[str],
    sample_rate: int = 16000,
    clip_duration: float = 3.0,
    transform: Callable | None = None,
)

Create an AudioDataset (torch Dataset) for batch audio processing.

This is a factory wrapper around a lazily-built torch.utils.data.Dataset subclass so the module imports without torch. The returned object IS a real Dataset instance.

Parameters:

Name	Type	Description	Default
filepaths	list[str]	List of audio file paths.	required
sample_rate	int	Target sample rate.	16000
clip_duration	float	Clip duration in seconds.	3.0
transform	Callable | None	Optional transform to apply.	None

create_dataloader

create_dataloader(
    filepaths: list[str],
    config: ModelConfig,
    transform: Callable | None = None,
) -> DataLoader

Create a DataLoader for batch processing.

Parameters:

Name	Type	Description	Default
filepaths	list[str]	List of audio file paths.	required
config	ModelConfig	Model configuration.	required
transform	Callable | None	Optional transform to apply.	None

Returns:

Type	Description
DataLoader	DataLoader instance.

get_model_class

get_model_class(name: str) -> type

Get a registered model class by name.

list_models

list_models() -> list[str]

List all registered model names.

register_model

register_model(name: str)

Decorator to register a model class.

batch_embed_files

batch_embed_files(
    input_dir: str,
    output_dir: str,
    model_path: str = "MIT/ast-finetuned-audioset-10-10-0.4593",
    *,
    layer: str = "last_hidden_state",
    normalize: bool = True,
    recursive: bool = True,
    max_workers: int = 1,
    continue_on_error: bool = True,
    on_progress: Callable[[int, int], None] | None = None,
) -> BatchResult

Extract AST embeddings for every audio file in a directory, saving one .npy per input file under output_dir.

The extractor (and its model) is loaded once and reused across files.

Parameters:

Name	Type	Description	Default
input_dir	str	Directory containing input audio files.	required
output_dir	str	Directory to write <stem>_embeddings.npy files to.	required
model_path	str	Path to model or HuggingFace identifier.	'MIT/ast-finetuned-audioset-10-10-0.4593'
layer	str	Layer to extract embeddings from.	'last_hidden_state'
normalize	bool	L2-normalize embeddings.	True
recursive	bool	Search subdirectories.	True
max_workers	int	Worker count (kept at 1 — the model is shared).	1
continue_on_error	bool	Collect per-file errors and keep going if True.	True
on_progress	Callable[[int, int], None] | None	Optional (completed, total) progress callback.	None

Returns:

Name	Type	Description
A	BatchResult	class:bioamla.batch.BatchResult summarizing the run.

Raises:

Type	Description
NotFoundError	If the input directory does not exist.
ModelError	On model-load or inference failure.

batch_predict_files

batch_predict_files(
    input_dir: str,
    model_path: str = "bioamla/scp-frogs",
    *,
    top_k: int = 5,
    min_confidence: float = 0.0,
    resample_freq: int = 16000,
    recursive: bool = True,
    max_workers: int = 1,
    continue_on_error: bool = True,
    on_progress: Callable[[int, int], None] | None = None,
) -> BatchResult

Run AST prediction over every audio file in a directory.

The model is loaded once and reused across files (sequential mode). Each file's structured prediction (top-k labels/scores) is collected into result.metadata["predictions"] so callers can write a structured predictions.json; result.output_files holds human-readable summaries.

Parameters:

Name	Type	Description	Default
input_dir	str	Directory containing input audio files.	required
model_path	str	Path to model or HuggingFace identifier.	'bioamla/scp-frogs'
top_k	int	Number of top predictions to keep per file.	5
min_confidence	float	Drop predictions below this probability.	0.0
resample_freq	int	Target sample rate.	16000
recursive	bool	Search subdirectories.	True
max_workers	int	Worker count (kept at 1 in practice — the model is shared).	1
continue_on_error	bool	Collect per-file errors and keep going if True.	True
on_progress	Callable[[int, int], None] | None	Optional (completed, total) progress callback.	None

Returns:

Name	Type	Description
A	BatchResult	class:bioamla.batch.BatchResult summarizing the run.

Raises:

Type	Description
NotFoundError	If the input directory does not exist.
ModelError	On model-load or inference failure.

batch_predict_segments

batch_predict_segments(
    input_dir: str,
    model_path: str = "bioamla/scp-frogs",
    *,
    segment_duration: int,
    overlap: int = 0,
    min_confidence: float = 0.0,
    resample_freq: int = 16000,
    recursive: bool = True,
    max_workers: int = 1,
    continue_on_error: bool = True,
    on_progress: Callable[[int, int], None] | None = None,
) -> BatchResult

Run segmented AST prediction over every audio file in a directory.

Each file is split into fixed-length (optionally overlapping) segments and classified per segment. The model is loaded once and reused across files (sequential mode). Per-segment rows (filepath, start_time, end_time, predicted_label, confidence) are collected into result.metadata["segments"] so callers can write a flat predictions.csv; result.output_files holds per-file summaries.

Parameters:

Name	Type	Description	Default
input_dir	str	Directory containing input audio files.	required
model_path	str	Path to model or HuggingFace identifier.	'bioamla/scp-frogs'
segment_duration	int	Duration of each segment in seconds.	required
overlap	int	Overlap between consecutive segments in seconds.	0
min_confidence	float	Drop segments whose prediction is below this probability.	0.0
resample_freq	int	Target sample rate.	16000
recursive	bool	Search subdirectories.	True
max_workers	int	Worker count (kept at 1 in practice — the model is shared).	1
continue_on_error	bool	Collect per-file errors and keep going if True.	True
on_progress	Callable[[int, int], None] | None	Optional (completed, total) progress callback.	None

Returns:

Name	Type	Description
A	BatchResult	class:bioamla.batch.BatchResult summarizing the run.

Raises:

Type	Description
NotFoundError	If the input directory does not exist.
ModelError	On model-load or inference failure.

get_current_device_index

get_current_device_index() -> int | None

Get the current CUDA device index, or None if not available.

get_device

get_device(prefer_cuda: bool = True) -> torch.device

Get the best available device for computation.

Parameters:

Name	Type	Description	Default
prefer_cuda	bool	If True (default), prefer CUDA if available.	True

Returns:

Type	Description
device	torch.device: The selected device (cuda or cpu)

get_device_count

get_device_count() -> int

Get the number of available CUDA devices.

get_device_info

get_device_info() -> dict

Get comprehensive information about available devices.

Returns:

Name	Type	Description
dict	dict	Device information including: - cuda_available: Whether CUDA is available - current_device: Current device index (if CUDA) - device_count: Number of CUDA devices - devices: List of device info dicts

get_device_name

get_device_name(device_index: int = 0) -> str | None

Get the name of a CUDA device.

Parameters:

Name	Type	Description	Default
device_index	int	The device index (default 0)	0

Returns:

Type	Description
str | None	Device name string, or None if not available

get_device_string

get_device_string(prefer_cuda: bool = True) -> str

Get the device as a string (for use with device_map="auto" etc.).

Parameters:

Name	Type	Description	Default
prefer_cuda	bool	If True (default), prefer CUDA if available.	True

Returns:

Name	Type	Description
str	str	"cuda" or "cpu"

is_cuda_available

is_cuda_available() -> bool

Check if CUDA is available.

move_to_device

move_to_device(
    model: Module, device: Union[str, device] | None = None
) -> nn.Module

Move a model to the specified device.

Parameters:

Name	Type	Description	Default
model	Module	PyTorch model to move	required
device	Union[str, device] | None	Target device. If None, uses get_device() to select.	None

Returns:

Type	Description
Module	The model on the target device

extract_embeddings

extract_embeddings(
    audio: str | ndarray,
    model_path: str = "MIT/ast-finetuned-audioset-10-10-0.4593",
    model_type: str = "ast",
    layer: str = "last_hidden_state",
    sample_rate: int | None = None,
    normalize: bool = True,
) -> np.ndarray

Extract embeddings from a single audio file or array.

For batch processing, use :class:EmbeddingExtractor directly.

Parameters:

Name	Type	Description	Default
audio	str | ndarray	Audio file path or numpy array.	required
model_path	str	Path to model or HuggingFace model ID.	'MIT/ast-finetuned-audioset-10-10-0.4593'
model_type	str	Model type ("ast").	'ast'
layer	str	Layer to extract from.	'last_hidden_state'
sample_rate	int | None	Sample rate (required if audio is an array).	None
normalize	bool	L2-normalize embeddings.	True

Returns:

Type	Description
ndarray	Embedding array of shape (n_segments, embedding_dim).

Raises:

Type	Description
ModelError	If extraction fails.

extract_embeddings_batch

extract_embeddings_batch(
    audio_files: list[str],
    model_path: str = "MIT/ast-finetuned-audioset-10-10-0.4593",
    model_type: str = "ast",
    output_path: str | None = None,
    output_format: str = "npy",
    aggregate: str = "mean",
    normalize: bool = True,
    progress_callback: Callable[[int, int], None]
    | None = None,
) -> BatchEmbeddingResult

Extract embeddings from multiple audio files, optionally saving them.

Parameters:

Name	Type	Description	Default
audio_files	list[str]	Audio file paths.	required
model_path	str	Path to model or HuggingFace model ID.	'MIT/ast-finetuned-audioset-10-10-0.4593'
model_type	str	Model type ("ast").	'ast'
output_path	str | None	Optional path to save embeddings.	None
output_format	str	Output format ("npy", "parquet", "csv").	'npy'
aggregate	str	Segment aggregation ("mean", "first", "all").	'mean'
normalize	bool	L2-normalize embeddings.	True
progress_callback	Callable[[int, int], None] | None	Optional (current, total) callback.	None

Returns:

Name	Type	Description
A	BatchEmbeddingResult	class:BatchEmbeddingResult.

get_ast_model_info

get_ast_model_info(model_path: str) -> dict[str, Any]

Return lightweight information about an AST model from its config.

Loads only the model config (not weights), so this is cheap.

Parameters:

Name	Type	Description	Default
model_path	str	Path to model or HuggingFace identifier.	required

Returns:

Type	Description
dict[str, Any]	A dict with path, model_type, num_classes, classes,
dict[str, Any]	has_more_classes, and hidden_size.

Raises:

Type	Description
ModelError	If the config cannot be loaded.

load_embeddings

load_embeddings(
    filepath: str, format: str | None = None
) -> tuple[np.ndarray, list[str]]

Load embeddings from disk.

Parameters:

Name	Type	Description	Default
filepath	str	Path to the embeddings file.	required
format	str | None	File format (auto-detected from the suffix if None).	None

Returns:

Type	Description
tuple[ndarray, list[str]]	(embeddings, filepaths).

Raises:

Type	Description
InvalidInputError	If the format is unknown.

save_embeddings

save_embeddings(
    embeddings: ndarray,
    filepaths: list[str],
    output_path: str,
    format: str = "npy",
    metadata: dict[str, Any] | None = None,
) -> str

Save embeddings to disk.

Parameters:

Name	Type	Description	Default
embeddings	ndarray	Embedding array.	required
filepaths	list[str]	Source file paths.	required
output_path	str	Output file path.	required
format	str	"npy", "npz", "parquet", or "csv".	'npy'
metadata	dict[str, Any] | None	Optional metadata (saved with the npz format).	None

Returns:

Type	Description
str	The path written to.

Raises:

Type	Description
InvalidInputError	If the format is unknown.

write_train_config

write_train_config(
    path: str | Path, force: bool = False
) -> Path

Write the AST training-config template to path.

Parameters:

Name	Type	Description	Default
path	str | Path	Destination file (parent directories are created).	required
force	bool	Overwrite an existing file instead of raising.	False

Returns:

Type	Description
Path	The written path.

Raises:

Type	Description
InvalidInputError	If path exists and force is False.

train_ast

train_ast(
    *,
    train_dataset: str,
    training_dir: str = ".",
    base_model: str = DEFAULT_BASE_MODEL,
    split: str = "train",
    category_label_column: str = "category",
    learning_rate: float = 5e-05,
    num_train_epochs: int = 1,
    per_device_train_batch_size: int = 8,
    eval_strategy: str = "epoch",
    save_strategy: str = "epoch",
    eval_steps: int = 1,
    save_steps: int = 1,
    load_best_model_at_end: bool = True,
    metric_for_best_model: str = "accuracy",
    logging_strategy: str = "steps",
    logging_steps: int = 100,
    report_to: str | list[str] = "tensorboard",
    fp16: bool = False,
    bf16: bool = False,
    gradient_accumulation_steps: int = 1,
    dataloader_num_workers: int = 4,
    torch_compile: bool = False,
    finetune_mode: str = "full",
    push_to_hub: bool = False,
    mlflow_tracking_uri: str | None = None,
    mlflow_experiment_name: str | None = None,
    mlflow_run_name: str | None = None,
    augmentation: AugmentationConfig | None = None,
    augment_multiplier: int = 1,
) -> TrainResult

Fine-tune an AST model on a custom dataset.

The remaining keyword arguments mirror transformers.TrainingArguments (learning rate, batch size, eval/save strategy, fp16/bf16, etc.).

Parameters:

Name	Type	Description	Default
train_dataset	str	A HuggingFace dataset id ("bioamla/scp-frogs" or "samuelstevens/BirdSet:HSN"), a local metadata CSV (with file and label columns), or a directory of class-named subdirectories.	required
training_dir	str	Output root; the best model is written to {training_dir}/best_model, runs to /runs, logs to /logs.	'.'
base_model	str	Pretrained AST checkpoint to fine-tune.	DEFAULT_BASE_MODEL
split	str	Split to use for single-split HuggingFace datasets.	'train'
category_label_column	str	Label column name for HF/CSV datasets.	'category'
augmentation	AugmentationConfig | None	On-the-fly training augmentation; None disables it.	None
augment_multiplier	int	Repeat the training split N times (with augmentation) to enlarge it; 1 means no duplication.	1

Returns:

Name	Type	Description
A	TrainResult	class:~bioamla.ml.ast_service.TrainResult with the best-model path,
	TrainResult	epoch count, and final eval accuracy/loss when available.

Raises:

Type	Description
TrainingError	On an unusable dataset/params or an empty training set.