openmusic_analysis/settings.py

from __future__ import annotations

import hashlib
import json
import os
from dataclasses import asdict, dataclass
from typing import Any


CLAP_MODEL_ID = "laion/clap-htsat-unfused"
CLAP_MODEL_REVISION = "8fa0f1c6d0433df6e97c127f64b2a1d6c0dcda8a"
BGE_MODEL_ID = "BAAI/bge-m3"
BGE_MODEL_REVISION = "5617a9f61b028005a4858fdac845db406aefb181"


def _env_bool(name: str, default: bool) -> bool:
    value = os.getenv(name)
    if value is None:
        return default
    return value.strip().lower() in {"1", "true", "yes", "on"}


@dataclass(frozen=True)
class GlobalAudioConfig:
    sample_rate: int = 48_000
    window_seconds: float = 10.0
    target_windows: int = 4
    coverage_start: float = 0.10
    coverage_end: float = 0.90
    minimum_audio_seconds: float = 1.0
    short_window_padding: str = "repeat_to_window_length"
    selection: str = "uniform_over_valid_start_range"
    aggregation: str = "l2_each_then_mean_then_l2"
    encoder_window_seconds: float = 10.0
    oversize_window_encoding: str = "uniform_10s_subwindows"
    encoder_subwindow_aggregation: str = "l2_each_then_mean_then_l2"
    decoding_version: str = "ffmpeg-mono-f32le-v1"
    resampling: str = "scipy-resample-poly-kaiser5-v1"
    feature_extractor: str = "transformers-clap-4.44.2-fixed-input"
    inference_batch_size: int = 4

    def __post_init__(self) -> None:
        if (
            self.sample_rate <= 0
            or self.window_seconds <= 0
            or self.target_windows <= 0
            or self.inference_batch_size <= 0
        ):
            raise ValueError("Global audio rates, window and count must be positive")
        if not 0 <= self.coverage_start <= self.coverage_end <= 1:
            raise ValueError("Global coverage must satisfy 0 <= start <= end <= 1")


@dataclass(frozen=True)
class TemporalAudioConfig:
    sample_rate: int = 48_000
    window_seconds: float = 10.0
    hop_seconds: float = 10.0
    max_segments: int = 24
    minimum_audio_seconds: float = 1.0
    short_window_padding: str = "repeat_to_window_length"
    segmentation: str = "sliding_with_tail_coverage_and_uniform_limit"
    encoder_window_seconds: float = 10.0
    oversize_window_encoding: str = "uniform_10s_subwindows"
    encoder_subwindow_aggregation: str = "l2_each_then_mean_then_l2"
    decoding_version: str = "ffmpeg-mono-f32le-v1"
    resampling: str = "scipy-resample-poly-kaiser5-v1"
    feature_extractor: str = "transformers-clap-4.44.2-fixed-input"
    inference_batch_size: int = 4

    def __post_init__(self) -> None:
        if (
            self.sample_rate <= 0
            or self.window_seconds <= 0
            or self.hop_seconds <= 0
            or self.max_segments <= 0
            or self.inference_batch_size <= 0
        ):
            raise ValueError("Temporal audio rates, windows, hop and limit must be positive")


@dataclass(frozen=True)
class LyricsConfig:
    max_chunk_tokens: int = 512
    batch_size: int = 8
    chunking: str = "section_first_then_token_split"
    aggregation: str = "token_weighted_mean_of_l2_chunks_then_l2"
    unicode_normalization: str = "NFKC"
    pooling: str = "bge-m3-cls"
    model_context_tokens: int = 8192

    def __post_init__(self) -> None:
        if self.max_chunk_tokens < 4 or self.max_chunk_tokens > self.model_context_tokens:
            raise ValueError("Lyrics chunk tokens must be between 4 and model context")
        if self.batch_size <= 0:
            raise ValueError("Lyrics batch size must be positive")


@dataclass(frozen=True)
class LimitConfig:
    max_upload_bytes: int = 100 * 1024 * 1024
    max_lyrics_characters: int = 100_000
    max_audio_seconds: float = 30 * 60
    request_timeout_seconds: float = 300.0
    decode_timeout_seconds: float = 120.0

    def __post_init__(self) -> None:
        if min(
            self.max_upload_bytes,
            self.max_lyrics_characters,
            self.max_audio_seconds,
            self.request_timeout_seconds,
            self.decode_timeout_seconds,
        ) <= 0:
            raise ValueError("All service limits must be positive")


@dataclass(frozen=True)
class Settings:
    device: str = "auto"
    inference_concurrency: int = 1
    ffmpeg_binary: str = "ffmpeg"
    ffprobe_binary: str = "ffprobe"
    eager_model_loading: bool = False
    global_audio: GlobalAudioConfig = GlobalAudioConfig()
    temporal_audio: TemporalAudioConfig = TemporalAudioConfig()
    lyrics: LyricsConfig = LyricsConfig()
    limits: LimitConfig = LimitConfig()

    def __post_init__(self) -> None:
        if self.inference_concurrency <= 0:
            raise ValueError("Inference concurrency must be positive")
        if self.global_audio.inference_batch_size != self.temporal_audio.inference_batch_size:
            raise ValueError("Global and temporal CLAP analyzers must share a batch size")

    @classmethod
    def from_env(cls) -> "Settings":
        clap_batch_size = int(os.getenv("OPENMUSIC_CLAP_BATCH_SIZE", "4"))
        global_audio = GlobalAudioConfig(
            window_seconds=float(os.getenv("OPENMUSIC_GLOBAL_WINDOW_SECONDS", "10")),
            target_windows=int(os.getenv("OPENMUSIC_GLOBAL_WINDOWS", "4")),
            inference_batch_size=clap_batch_size,
        )
        temporal_audio = TemporalAudioConfig(
            window_seconds=float(os.getenv("OPENMUSIC_TEMPORAL_WINDOW_SECONDS", "10")),
            hop_seconds=float(os.getenv("OPENMUSIC_TEMPORAL_HOP_SECONDS", "10")),
            max_segments=int(os.getenv("OPENMUSIC_TEMPORAL_MAX_SEGMENTS", "24")),
            inference_batch_size=clap_batch_size,
        )
        lyrics = LyricsConfig(
            max_chunk_tokens=int(os.getenv("OPENMUSIC_LYRICS_CHUNK_TOKENS", "512")),
            batch_size=int(os.getenv("OPENMUSIC_LYRICS_BATCH_SIZE", "8")),
        )
        limits = LimitConfig(
            max_upload_bytes=int(os.getenv("OPENMUSIC_MAX_UPLOAD_BYTES", str(100 * 1024 * 1024))),
            max_lyrics_characters=int(os.getenv("OPENMUSIC_MAX_LYRICS_CHARACTERS", "100000")),
            max_audio_seconds=float(os.getenv("OPENMUSIC_MAX_AUDIO_SECONDS", "1800")),
            request_timeout_seconds=float(os.getenv("OPENMUSIC_REQUEST_TIMEOUT_SECONDS", "300")),
            decode_timeout_seconds=float(os.getenv("OPENMUSIC_DECODE_TIMEOUT_SECONDS", "120")),
        )
        return cls(
            device=os.getenv("OPENMUSIC_DEVICE", "auto"),
            inference_concurrency=max(1, int(os.getenv("OPENMUSIC_INFERENCE_CONCURRENCY", "1"))),
            ffmpeg_binary=os.getenv("OPENMUSIC_FFMPEG", "ffmpeg"),
            ffprobe_binary=os.getenv("OPENMUSIC_FFPROBE", "ffprobe"),
            eager_model_loading=_env_bool("OPENMUSIC_EAGER_MODELS", False),
            global_audio=global_audio,
            temporal_audio=temporal_audio,
            lyrics=lyrics,
            limits=limits,
        )


def config_dict(config: Any) -> dict[str, Any]:
    return asdict(config)


def preprocessing_version(base: str, config: Any) -> str:
    """Tie a preprocessing version to every result-affecting configuration value."""
    payload = json.dumps(config_dict(config), sort_keys=True, separators=(",", ":"))
    digest = hashlib.sha256(payload.encode("utf-8")).hexdigest()[:12]
    return f"{base}.{digest}"