extraction.py

# MIT License
# 
# Copyright (c) 2026 audio-embeddings contributors
# 
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# 
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from __future__ import annotations

from dataclasses import dataclass
from typing import Callable

import torch


@dataclass(frozen=True)
class ExtractionPreset:
    name: str
    overlap: float
    num_phases: int


PRESETS = {
    preset.name: preset
    for preset in (
        ExtractionPreset("native_one_phase", 0.0, 1),
        ExtractionPreset("overlap25_one_phase", 0.25, 1),
        ExtractionPreset("overlap50_one_phase", 0.5, 1),
        ExtractionPreset("native_two_phase", 0.0, 2),
        ExtractionPreset("overlap25_two_phase", 0.25, 2),
        ExtractionPreset("overlap50_two_phase", 0.5, 2),
    )
}


def get_preset(name: str) -> ExtractionPreset:
    try:
        return PRESETS[name]
    except KeyError as error:
        raise ValueError(
            f"Unknown extraction preset {name!r}. Expected one of {sorted(PRESETS)}"
        ) from error


def _window_starts(length: int, window: int, overlap: float) -> list[int]:
    if length <= window:
        return [0]
    if not 0.0 <= overlap < 1.0:
        raise ValueError(f"overlap must be in [0, 1), got {overlap}")
    stride = max(1, round(window * (1.0 - overlap)))
    if overlap == 0.0:
        return list(range(0, length, stride))
    starts = list(range(0, length - window + 1, stride))
    final_start = length - window
    if starts[-1] != final_start:
        starts.append(final_start)
    return starts


def _positive_triangular_weights(
    length: int,
    *,
    device: torch.device,
    dtype: torch.dtype,
) -> torch.Tensor:
    positions = torch.arange(length, device=device, dtype=dtype)
    return 1.0 - torch.abs((2.0 * positions) - (length - 1)) / (length + 1)


def fuse_context_windows(
    tokens: torch.Tensor,
    *,
    max_context_tokens: int,
    overlap: float,
    encode_window: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
) -> torch.Tensor:
    """Contextualize a [frequency, time, dim] grid and fuse repeated tokens."""

    if tokens.ndim != 3:
        raise ValueError(
            f"Expected token grid [frequency, time, dim], got {tuple(tokens.shape)}"
        )
    frequency, time, dimension = tokens.shape
    max_time = max_context_tokens // frequency
    if max_time <= 0:
        raise ValueError(
            f"Encoder context {max_context_tokens} cannot hold {frequency} frequency tokens"
        )
    if time == 0:
        raise ValueError("Token grid has no time steps")

    accumulator = torch.zeros_like(tokens)
    denominator = torch.zeros(time, device=tokens.device, dtype=tokens.dtype)
    for start in _window_starts(time, max_time, overlap):
        end = min(time, start + max_time)
        width = end - start
        window = tokens[:, start:end, :]
        flattened = window.reshape(1, frequency * width, dimension)
        position_ids = torch.arange(
            frequency * width,
            device=tokens.device,
        )
        encoded = encode_window(flattened, position_ids)
        if encoded.shape != flattened.shape:
            raise ValueError(
                "Encoder changed the token-grid shape: "
                f"expected {tuple(flattened.shape)}, got {tuple(encoded.shape)}"
            )
        encoded_grid = encoded.reshape(frequency, width, dimension)
        weights = _positive_triangular_weights(
            width,
            device=tokens.device,
            dtype=tokens.dtype,
        )
        accumulator[:, start:end, :] += encoded_grid * weights.view(1, -1, 1)
        denominator[start:end] += weights

    if torch.any(denominator <= 0):
        raise RuntimeError("At least one token received zero overlap weight")
    fused = accumulator / denominator.view(1, -1, 1)
    return fused.mean(dim=0)


def merge_phases(
    phases: list[tuple[torch.Tensor, torch.Tensor]],
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    if not phases:
        raise ValueError("At least one embedding phase is required")
    embedding_dim = phases[0][0].shape[-1]
    for embeddings, timestamps in phases:
        if embeddings.ndim != 2 or embeddings.shape[-1] != embedding_dim:
            raise ValueError("All phase embeddings must have shape [time, dimension]")
        if timestamps.ndim != 1 or timestamps.shape[0] != embeddings.shape[0]:
            raise ValueError("Each phase needs one timestamp per embedding")

    merged_embeddings = torch.cat([phase[0] for phase in phases], dim=0)
    merged_timestamps = torch.cat([phase[1] for phase in phases], dim=0)
    order = torch.argsort(merged_timestamps, stable=True)
    scene = torch.stack([embeddings.mean(dim=0) for embeddings, _ in phases]).mean(
        dim=0
    )
    return merged_embeddings[order], merged_timestamps[order], scene