music-assistant-server

1"""
2Precompute CLAP text embeddings for SCALAR_PROMPT_PAIRS and ship as .npz.
3
4Run once whenever the prompts in clap_prompts.py change (and bump the
5provider's analysis_version alongside). The resulting artifact is loaded
6by the sonic_analysis provider so the audio-only path can score scalars
7without downloading the GPT2 text encoder.
8
9Usage:
10    python scripts/precompute_clap_prompt_embeddings.py
11"""
12
13from __future__ import annotations
14
15import sys
16
17# ruff: noqa: T201
18from music_assistant.providers.sonic_analysis.clap_prompts import (
19    PRECOMPUTED_EMBEDDINGS_PATH,
20    SCALAR_PROMPT_PAIRS,
21    compute_prompt_embeddings,
22    hash_scalar_prompt_pairs,
23    save_precomputed_prompt_embeddings,
24)
25from music_assistant.providers.sonic_analysis.vendored_clap import CLAP
26
27
28def main() -> int:
29    """Load CLAP, embed SCALAR_PROMPT_PAIRS, and write the .npz artifact."""
30    print("Loading CLAP model with text encoder (this triggers GPT2 download on first run)...")
31    model = CLAP(version="2023", use_cuda=False)
32
33    print(
34        f"Embedding {len(SCALAR_PROMPT_PAIRS)} prompt pairs ({2 * len(SCALAR_PROMPT_PAIRS)} strings)..."
35    )
36    embeddings = compute_prompt_embeddings(model, SCALAR_PROMPT_PAIRS)
37    prompts_hash = hash_scalar_prompt_pairs(SCALAR_PROMPT_PAIRS)
38
39    print(f"Embeddings shape: {embeddings.shape}, dtype: {embeddings.dtype}")
40    print(f"Prompts hash: {prompts_hash}")
41
42    PRECOMPUTED_EMBEDDINGS_PATH.parent.mkdir(parents=True, exist_ok=True)
43    save_precomputed_prompt_embeddings(PRECOMPUTED_EMBEDDINGS_PATH, embeddings, prompts_hash)
44
45    size_kb = PRECOMPUTED_EMBEDDINGS_PATH.stat().st_size / 1024
46    print(f"Wrote {PRECOMPUTED_EMBEDDINGS_PATH} ({size_kb:.1f} KB)")
47    return 0
48
49
50if __name__ == "__main__":
51    sys.exit(main())
52