sync_pilot/config.py

from __future__ import annotations

import logging
import os
from pathlib import Path
from typing import TYPE_CHECKING, Any

from dotenv import load_dotenv

if TYPE_CHECKING:
    # Only loaded by type checkers; the runtime import is deferred to
    # ``_detect_device``/``get_default_dtype`` so dashboard-only deployments
    # (e.g. the HF Space) can import this module without the heavy ML stack.
    import torch

# MPS has incomplete op coverage; this lets unsupported ops fall back to CPU
# instead of raising NotImplementedError mid-inference.
os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")

PACKAGE_ROOT: Path = Path(__file__).resolve().parent
PROJECT_ROOT: Path = PACKAGE_ROOT.parent

load_dotenv(PROJECT_ROOT / ".env", override=False)

PIPELINE_VERSION: str = "0.0.1"

DATA_DIR: Path = Path(os.getenv("SYNC_PILOT_DATA_DIR", PROJECT_ROOT / "data"))
AUDIO_DIR: Path = DATA_DIR / "audio"
CACHE_DIR: Path = Path(os.getenv("SYNC_PILOT_CACHE_DIR", DATA_DIR / "cache"))
OUTPUT_DIR: Path = Path(os.getenv("SYNC_PILOT_OUTPUT_DIR", DATA_DIR / "outputs"))
EMBEDDINGS_DIR: Path = DATA_DIR / "embeddings"
LOGS_DIR: Path = DATA_DIR / "logs"

HF_CACHE_DIR: Path = Path(os.getenv("HF_HOME", Path.home() / ".cache" / "huggingface"))

KB_PATH: Path = PROJECT_ROOT / "kb" / "music_models.json"

HF_TOKEN: str | None = os.getenv("HF_TOKEN") or None
OPENROUTER_API_KEY: str | None = os.getenv("OPENROUTER_API_KEY") or None
# Discogs personal access token. When set, ``groundtruth/sources.fetch_discogs``
# uses the authenticated v2 API (60 req/min) instead of webpage scraping
# (which Discogs 403s for generic User-Agents). When unset, Discogs coverage
# is silently skipped — a one-time warning is logged.
DISCOGS_TOKEN: str | None = os.getenv("DISCOGS_TOKEN") or None

# LLM config — used by sync_pilot.description to synthesize per-track prose
# from the structured tag layers. Default is DeepSeek (cheap, fluent EN, copes
# with Turkish loanwords). Override LLM_MODEL to A/B with anthropic/claude-*
# or openai/gpt-* variants without code changes. Pricing (USD per 1M tokens)
# is model-dependent; we record per-call usage in the batch summary so cost
# can be re-estimated if the rate sheet drifts.
LLM_BASE_URL: str = os.getenv("LLM_BASE_URL", "https://openrouter.ai/api/v1")
LLM_MODEL: str = os.getenv("LLM_MODEL", "deepseek/deepseek-chat")
# Content-safety classification defaults to a Sonnet-class Claude: in the
# 2026-06 A/B it caught both true-explicit tracks (050, 070) while deepseek-chat
# missed the profanity track 050, at comparable cost/latency. Theme extraction
# and all other LLM calls stay on LLM_MODEL (deepseek-chat) — Claude
# under-generated themes (empty on several tracks) and R1 was too slow.
CONTENT_SAFETY_LLM_MODEL: str = os.getenv(
    "SYNC_PILOT_CONTENT_SAFETY_MODEL", "anthropic/claude-sonnet-4.6"
)
# Live Sync rerank is the user-facing, judgment-heavy, Turkish-brief step — the
# one place a frontier model most visibly beats deepseek-chat (ranking quality +
# native structured output, no ```json-fence fragility). Kept env-overridable so
# it can be A/B'd against deepseek without code changes.
SYNC_RERANK_LLM_MODEL: str = os.getenv(
    # Haiku for the live rerank: ~2-3x snappier than Sonnet at comparable ranking
    # quality, cheaper, higher rate limits — the right trade for a shared sync demo
    # where the LLM call is ~95% of the latency. Override the env to A/B vs Sonnet.
    "SYNC_PILOT_SYNC_RERANK_MODEL", "anthropic/claude-haiku-4.5"
)
# Embedding model for Live Sync retrieval ONLY. Decoupled from the lyla text
# classifier (which stays on multilingual-e5-base, the encoder its heads were
# trained on). e5-large is a drop-in upgrade — same query:/passage: prefix
# scheme, 1024d. Build (scripts/build_sync_index) and query (sync_recommend)
# read this together, so the index and query stay dimension-consistent.
SYNC_EMBED_MODEL: str = os.getenv(
    "SYNC_PILOT_SYNC_EMBED_MODEL", "intfloat/multilingual-e5-large"
)
# Track descriptions are user-facing Turkish prose that must respect the GT
# grounding (the specialist flagged deepseek mislabelling genre — "rock şarkısına
# halk türküsü" — and producing prose off that). A frontier model follows the
# grounding block far more reliably. Env-overridable to A/B against deepseek.
DESCRIPTION_LLM_MODEL: str = os.getenv(
    "SYNC_PILOT_DESCRIPTION_MODEL", "anthropic/claude-sonnet-4.6"
)
# Sync brief (the per-track placement recommendation shown in the showcase) and
# the sync search keywords are user-facing judgment tasks — same rationale as
# descriptions, so they move to the frontier model too. Theme extraction stays on
# LLM_MODEL (deepseek): Claude under-generated themes (documented 2026-06).
SYNC_BRIEF_LLM_MODEL: str = os.getenv("SYNC_PILOT_SYNC_BRIEF_MODEL", "anthropic/claude-sonnet-4.6")
KEYWORDS_LLM_MODEL: str = os.getenv("SYNC_PILOT_KEYWORDS_MODEL", "anthropic/claude-sonnet-4.6")
LLM_TEMPERATURE: float = float(os.getenv("LLM_TEMPERATURE", "0.4"))
# DeepSeek-chat current OpenRouter list price (May 2026): roughly $0.27/1M in,
# $1.10/1M out. Treat these as defaults for the cost estimate column in the
# batch summary; bump from the env if the user is benching a different model.
LLM_INPUT_PRICE_PER_M: float = float(os.getenv("LLM_INPUT_PRICE_PER_M", "0.27"))
LLM_OUTPUT_PRICE_PER_M: float = float(os.getenv("LLM_OUTPUT_PRICE_PER_M", "1.10"))


def _detect_device() -> "torch.device":
    import torch

    override = os.getenv("SYNC_PILOT_DEVICE", "").strip().lower()
    if override in {"mps", "cuda", "cpu"}:
        return torch.device(override)
    if torch.backends.mps.is_available():
        return torch.device("mps")
    if torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")


# ``DEVICE`` is computed lazily so importing this module never triggers a
# ``torch`` import. Pipeline code goes through ``get_device()``; the cached
# value is reused for the rest of the process lifetime.
_DEVICE: "torch.device | None" = None


def get_device() -> "torch.device":
    global _DEVICE
    if _DEVICE is None:
        _DEVICE = _detect_device()
    return _DEVICE


def get_default_dtype() -> "torch.dtype":
    import torch

    return _dtype_by_device().get(get_device().type, torch.float32)


def _dtype_by_device() -> dict[str, "torch.dtype"]:
    import torch

    # SYNC_PILOT_FORCE_FP32 forces float32 on every device. Used on the live
    # backend's CUDA GPU so (a) inference matches the MPS-developed pipeline
    # bit-for-bit-ish (reproducibility) and (b) we avoid fp16-input vs fp32-weight
    # mismatches — MAEST loads its weights in fp32, so casting inputs to fp16
    # (the cuda default below) raised "Input type (Half) and bias type (float)".
    if os.getenv("SYNC_PILOT_FORCE_FP32") == "1":
        return {"mps": torch.float32, "cuda": torch.float32, "cpu": torch.float32}
    # bfloat16 is unsupported on MPS — keep dtype map device-aware.
    return {
        "mps": torch.float32,
        "cuda": torch.float16,
        "cpu": torch.float32,
    }


def __getattr__(name: str) -> Any:
    """Module-level lazy attribute access for the heavy-ML constants.

    Pipeline modules (``tagging.py``, ``clap_tagging.py``, ``transcription.py``)
    still reference ``config.DEVICE`` and ``config.DTYPE_BY_DEVICE`` directly.
    Defining them at module scope would re-pull ``torch`` on every import,
    which is what we moved away from so dashboard-only deployments stay
    light. PEP 562 ``__getattr__`` lets us keep the names available while
    deferring the ``torch`` import to first access.
    """
    if name == "DEVICE":
        return get_device()
    if name == "DTYPE_BY_DEVICE":
        return _dtype_by_device()
    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")


# Source of truth for model selection is sync_pilot/kb/music_models.json.
# When the music-models-scout updates the KB, mirror the winning hf_id here.
MODELS: dict[str, dict[str, Any]] = {
    "tagging": {
        "hf_id": "mtg-upf/discogs-maest-30s-pw-129e",
        "library": "transformers",
        "trust_remote_code": True,
        "sample_rate_hz": 16000,
        "segment_seconds": 30,
        "license": "cc-by-nc-sa-4.0",
        "license_class": "non-commercial-research",
    },
    "zero_shot_tagging": {
        # KB lists "laion/larger_clap_music" but that checkpoint ships with a
        # broken logit_scale (~1.0 instead of trained ~14-38), which collapses
        # softmax to ~uniform across prompts and makes zero-shot scoring
        # useless. We use "laion/clap-htsat-unfused" instead — also music-
        # trained (HTSAT audio backbone, MusicCaps-style captions), Apache-2.0,
        # 48 kHz, and produces well-calibrated zero-shot logits in the
        # transformers ClapModel API. Verified empirically on 038_dert_olur
        # (Anadolu rock track) where it cleanly separated Anatolian rock from
        # heavy metal and Turkish folk prompts.
        "hf_id": "laion/clap-htsat-unfused",
        "library": "transformers",
        "trust_remote_code": False,
        "sample_rate_hz": 48000,
        "license": "apache-2.0",
        "license_class": "permissive-commercial-ok",
    },
    "lyrics_asr": {
        # Whisper-large-v3-turbo: modern Apache-2.0 OpenAI checkpoint, ~5x
        # faster than large-v3 with negligible quality loss on high-resource
        # languages (Turkish is well-covered). ~800 MB on disk. Override the
        # hf_id with SYNC_PILOT_LYRICS_ASR_MODEL to A/B against
        # openai/whisper-large-v3 (slower, slightly more accurate) or
        # openai/whisper-medium (smaller, faster). chunk_length_s=30 matches
        # Whisper's native receptive field; stride_length_s=5 gives the
        # transformers ASR pipeline overlap to stitch long-form outputs
        # without repeating phrases at chunk boundaries.
        "hf_id": os.getenv(
            "SYNC_PILOT_LYRICS_ASR_MODEL", "openai/whisper-large-v3-turbo"
        ),
        "library": "transformers",
        "trust_remote_code": False,
        "sample_rate_hz": 16000,
        "chunk_length_s": 30,
        "stride_length_s": 5,
        "language": "turkish",
        "task": "transcribe",
        "license": "apache-2.0",
        "license_class": "permissive-commercial-ok",
    },
    "embeddings": {
        # TODO(music-models-scout): MERT-v1-95M / 330M candidate
        "hf_id": None,
    },
    "structural_segmentation": {
        # TODO(music-models-scout)
        "hf_id": None,
    },
    "audio_captioning": {
        # TODO(music-models-scout)
        "hf_id": None,
    },
    "stem_separation": {
        # TODO(music-models-scout) — Phase 2
        "hf_id": None,
    },
}

# Canonical sample rate for normalized ingest. Stage modules resample from this
# to their model-specific rate (see MODELS[stage]["sample_rate_hz"]).
INGEST_SAMPLE_RATE_HZ: int = 44100
INGEST_CHANNELS: int = 1


def ensure_dirs() -> None:
    for d in (AUDIO_DIR, CACHE_DIR, OUTPUT_DIR, EMBEDDINGS_DIR, LOGS_DIR):
        d.mkdir(parents=True, exist_ok=True)


def stage_cache_dir(stage: str) -> Path:
    p = CACHE_DIR / stage
    p.mkdir(parents=True, exist_ok=True)
    return p


def setup_logging(level: int = logging.INFO) -> None:
    logging.basicConfig(
        level=level,
        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
        datefmt="%Y-%m-%dT%H:%M:%S",
    )