Upload folder using huggingface_hub

README.md

@@ -1,13 +1,38 @@
 ---
-title: Steerable Retrieval
-emoji: 🔥
-colorFrom: gray
-colorTo: pink
+title: Sparse Steerable Retrieval
+emoji: 🎚️
+colorFrom: purple
+colorTo: indigo
 sdk: gradio
-sdk_version: 6.20.0
-python_version: '3.13'
+sdk_version: 4.44.0
 app_file: app.py
 pinned: false
+license: mit
+short_description: Steer dense music retrieval along a free-form concept
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Sparse Steerable Retrieval — live demo
+
+Open-vocabulary concept control for dense music retrieval. Pick a seed track, type a
+free-form concept, and steer: the query embedding is edited along the concept axis via
+sparse inversion in a trained BatchTopK SAE (MuQ-MuLan / music4all), then nearest
+neighbours are retrieved. No audio is hosted — playback is via Spotify embeds.
+
+Runs on **ZeroGPU**: the corpus + MuQ weights are fetched at startup; the model is built
+on the GPU lazily inside the first request and cached.
+
+## Configuration (Space variables / secrets)
+
+- `SSR_CHECKPOINT` — HF model repo id, e.g. `Pliploop/steerable-retrieval-sae` (variable)
+- `SSR_L0` — subfolder to load, e.g. `L0-20` (variable)
+- `SSR_CORPUS_REPO` — HF **dataset** repo id with `corpus.npz` + `meta.json` (variable)
+- `SSR_DEVICE` — `cuda` (variable)
+- `HF_TOKEN` — a read token (secret) so the Space can read the private corpus dataset
+
+## Layout
+
+- `app.py` — the Gradio ZeroGPU app.
+- `core.py` — engine: load model + corpus, steer + retrieve.
+- `steerable_retrieval/` — vendored package (steering API + SAE + MuQ encoder).
+
+This Space is published by `hf/push_all.py` in the project repo.

app.py

@@ -0,0 +1,142 @@
+"""Sparse Steerable Retrieval — live demo (Hugging Face ZeroGPU Space).
+
+Pick a seed track, type a free-form concept, and steer: the query embedding is edited
+along the concept axis via sparse inversion in a trained SAE, then nearest neighbours are
+retrieved over a music4all corpus. No audio is hosted — playback is via Spotify embeds.
+
+ZeroGPU: the corpus + MuQ weights are fetched at startup (CPU); the model is built on the
+GPU lazily inside the first @spaces.GPU call and cached for subsequent requests.
+"""
+import os
+import random
+
+import gradio as gr
+
+# `spaces` exists only on HF ZeroGPU; shim to a no-op decorator elsewhere.
+try:
+    import spaces
+except Exception:  # pragma: no cover
+    class _Spaces:
+        def GPU(self, *a, **k):
+            def deco(fn):
+                return fn
+            return deco if not (a and callable(a[0])) else a[0]
+    spaces = _Spaces()
+
+from core import DemoEngine, load_corpus
+
+K = 8
+EXAMPLE_CONCEPTS = ["piano", "distorted guitar", "female vocals", "aggressive and intense",
+                    "dreamy and ethereal", "warm and intimate", "hip hop beats"]
+
+# --- startup (CPU): pre-cache weights so the first GPU call stays within budget -------- #
+for _repo in ("OpenMuQ/MuQ-MuLan-large", "OpenMuQ/MuQ-large-msd-iter"):
+    try:
+        from huggingface_hub import snapshot_download
+        snapshot_download(_repo)
+    except Exception:
+        pass
+
+CORPUS = load_corpus()
+print(f"corpus loaded: {len(CORPUS[1])} tracks")
+
+_ENGINE = None
+
+
+def ensure_engine():
+    global _ENGINE
+    if _ENGINE is None:
+        _ENGINE = DemoEngine(corpus=CORPUS, dev="cuda")
+    return _ENGINE
+
+
+# --- rendering ------------------------------------------------------------------------ #
+def _spotify(sid, height=80):
+    return (
+        f'<iframe style="border-radius:12px" src="https://open.spotify.com/embed/track/{sid}?theme=0" '
+        f'width="100%" height="{height}" frameBorder="0" loading="lazy" '
+        f'allow="autoplay; clipboard-write; encrypted-media; fullscreen; picture-in-picture"></iframe>'
+    )
+
+
+def _meta(track_id):
+    m = CORPUS[2].get(track_id, {})
+    return m.get("title", ""), m.get("artist", ""), m.get("genre", ""), m.get("spotify", "")
+
+
+def seed_card(track_id):
+    t, a, g, sid = _meta(track_id)
+    return f'<div class="ssr-seed"><div class="ssr-title">{t}</div><div class="ssr-sub">{a} · {g}</div>{_spotify(sid, 152)}</div>'
+
+
+def results_html(results):
+    cards = []
+    for i, r in enumerate(results, 1):
+        cards.append(
+            f'<div class="ssr-card"><div class="ssr-row"><span class="ssr-rank">{i}</span>'
+            f'<span class="ssr-title">{r["title"]}</span> <span class="ssr-sub">{r["artist"]}</span>'
+            f'<span class="ssr-aff">{r["affinity"]:.2f}</span></div>{_spotify(r["spotify"])}</div>'
+        )
+    return '<div class="ssr-results">' + "".join(cards) + "</div>"
+
+
+def random_seed():
+    tid = random.choice(CORPUS[1])
+    return tid, seed_card(tid)
+
+
+@spaces.GPU(duration=120)
+def steer(seed_track_id, concept, alpha):
+    if not seed_track_id:
+        return "<em>Pick a seed track first.</em>", ""
+    if not concept or not concept.strip():
+        return "<em>Type a concept to steer toward.</em>", ""
+    eng = ensure_engine()
+    concept = concept.strip()
+    results = eng.steer_and_retrieve(seed_track_id, concept, alpha=float(alpha), k=K)
+    support = len(eng._slider(concept))
+    verb = "amplifying" if alpha >= 0 else "suppressing"
+    note = f"{verb} <b>{concept}</b> (α={float(alpha):+.1f}) · concept support = {support} SAE features"
+    return note, results_html(results)
+
+
+CSS = """
+.gradio-container {max-width: 1100px !important}
+.ssr-seed .ssr-title, .ssr-card .ssr-title {font-weight:600}
+.ssr-sub {color:#8a8a8a; font-size:13px}
+.ssr-results {display:flex; flex-direction:column; gap:10px}
+.ssr-card {padding:8px 10px; border-radius:14px; background:rgba(123,63,242,0.05)}
+.ssr-row {display:flex; align-items:center; gap:8px; margin-bottom:6px}
+.ssr-rank {display:inline-flex; align-items:center; justify-content:center; width:20px; height:20px;
+  border-radius:999px; background:#7B3FF2; color:#fff; font-size:11px; font-weight:600; flex:none}
+.ssr-aff {margin-left:auto; color:#aaa; font-size:12px}
+"""
+
+with gr.Blocks(title="Sparse Steerable Retrieval", css=CSS, theme=gr.themes.Soft(primary_hue="purple")) as demo:
+    gr.Markdown(
+        "# 🎚️ Sparse Steerable Retrieval\n"
+        "Steer a seed track along a **free-form concept**, then retrieve. A sparse autoencoder "
+        "factorises the MuQ-MuLan embedding into concept features; sparse inversion finds the ones "
+        "your concept maps to, so they can be amplified (**α > 0**) or suppressed (**α < 0**) before "
+        "nearest-neighbour search over ~109k music4all tracks. *No audio is hosted — playback via Spotify.*"
+    )
+    seed_state = gr.State()
+    with gr.Row():
+        with gr.Column(scale=1):
+            seed_html = gr.HTML()
+            random_btn = gr.Button("🎲 Random seed track", variant="secondary")
+            concept = gr.Textbox(label="Concept", value="piano", placeholder="e.g. distorted guitar, dreamy, aggressive")
+            gr.Examples(EXAMPLE_CONCEPTS, inputs=concept, label="Try a concept")
+            alpha = gr.Slider(-3.0, 3.0, value=1.5, step=0.1, label="α  (– suppress · + amplify)")
+            go = gr.Button("Steer & retrieve", variant="primary")
+            note = gr.HTML()
+        with gr.Column(scale=1):
+            results_out = gr.HTML()
+
+    random_btn.click(random_seed, outputs=[seed_state, seed_html])
+    go.click(steer, inputs=[seed_state, concept, alpha], outputs=[note, results_out])
+    demo.load(random_seed, outputs=[seed_state, seed_html])
+
+
+if __name__ == "__main__":
+    demo.launch()

core.py

@@ -0,0 +1,98 @@
+"""Shared engine for the demo Space and website example generation.
+
+Loads a trained BatchTopK SAE + MuQ text tower and a music4all corpus, then steers a
+seed track toward a free-form concept and retrieves nearest neighbours. Corpus load
+(CPU, cheap) is separated from model load (GPU, MuQ) so the Space can fetch the corpus at
+startup and build the model lazily inside a ZeroGPU call.
+
+Configuration (env):
+  SSR_CHECKPOINT   local .ckpt path, or a HF model repo id (default: local L0=20 run)
+  SSR_L0           subfolder in the model repo to load, e.g. "L0-20" (Hub repos only)
+  SSR_CORPUS_REPO  HF dataset repo id holding corpus.npz + meta.json (overrides local dir)
+  SSR_CORPUS_DIR   local corpus dir (default: demo/corpus)
+  SSR_DEVICE       "cuda" / "cpu" (auto-detected)
+"""
+from __future__ import annotations
+
+import json
+import os
+from functools import lru_cache
+
+import numpy as np
+import torch
+
+from steerable_retrieval.steer import Slider
+from steerable_retrieval.steer.loading import load_steerable_sae
+
+HERE = os.path.dirname(__file__)
+CORPUS_DIR = os.environ.get("SSR_CORPUS_DIR", os.path.join(HERE, "corpus"))
+CORPUS_REPO = os.environ.get("SSR_CORPUS_REPO")  # HF dataset repo id, if hosted
+CHECKPOINT = os.environ.get("SSR_CHECKPOINT", os.path.join(HERE, os.pardir, "logs/xps/66a0caf9/checkpoints/last.ckpt"))
+L0_SUBFOLDER = os.environ.get("SSR_L0")  # e.g. "L0-20" when CHECKPOINT is a Hub repo
+CONFIG = os.environ.get("SSR_CONFIG")
+
+
+def device() -> str:
+    return os.environ.get("SSR_DEVICE") or ("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def load_corpus():
+    """Return (embeddings[np], track_ids[list], meta[dict]). Fetches from the HF dataset
+    repo if SSR_CORPUS_REPO is set, else reads the local corpus dir. CPU only."""
+    if CORPUS_REPO:
+        from huggingface_hub import hf_hub_download
+
+        npz_path = hf_hub_download(CORPUS_REPO, filename="corpus.npz", repo_type="dataset")
+        meta_path = hf_hub_download(CORPUS_REPO, filename="meta.json", repo_type="dataset")
+    else:
+        npz_path = os.path.join(CORPUS_DIR, "corpus.npz")
+        meta_path = os.path.join(CORPUS_DIR, "meta.json")
+    npz = np.load(npz_path, allow_pickle=True)
+    with open(meta_path) as fh:
+        meta = json.load(fh)
+    return npz["embeddings"].astype(np.float32), [str(t) for t in npz["track_ids"].tolist()], meta
+
+
+class DemoEngine:
+    def __init__(self, corpus=None, dev: str | None = None):
+        self.device = dev or device()
+        embs, ids, meta = corpus if corpus is not None else load_corpus()
+        self.embeddings = torch.from_numpy(embs).to(self.device)
+        self.track_ids = ids
+        self.id_to_idx = {t: i for i, t in enumerate(ids)}
+        self.meta = meta
+        self.model = load_steerable_sae(CHECKPOINT, device=self.device, subfolder=L0_SUBFOLDER, config_path=CONFIG)
+
+    @lru_cache(maxsize=64)
+    def _slider(self, concept: str) -> Slider:
+        return Slider(concept, model=self.model, method="adam")
+
+    def track_meta(self, track_id: str) -> dict:
+        m = dict(self.meta.get(track_id, {}))
+        m["track_id"] = track_id
+        return m
+
+    def seed_embedding(self, track_id: str) -> torch.Tensor:
+        return self.embeddings[self.id_to_idx[track_id]]
+
+    def steer_and_retrieve(self, seed_track_id: str, concept: str, alpha: float = 1.0, k: int = 8) -> list[dict]:
+        slider = self._slider(concept)
+        z = self.seed_embedding(seed_track_id)
+        seed_idx = self.id_to_idx[seed_track_id]
+        idx, scores = slider.retrieve(z, self.embeddings, alpha=alpha, k=k, exclude_idx=seed_idx)
+        out = []
+        for i, s in zip(idx.tolist(), scores.tolist()):
+            m = self.track_meta(self.track_ids[i])
+            m["affinity"] = float(s)
+            out.append(m)
+        return out
+
+
+_ENGINE: DemoEngine | None = None
+
+
+def get_engine(corpus=None) -> DemoEngine:
+    global _ENGINE
+    if _ENGINE is None:
+        _ENGINE = DemoEngine(corpus=corpus)
+    return _ENGINE

requirements.txt

@@ -0,0 +1,15 @@
+# Hugging Face ZeroGPU Space runtime deps. The steerable_retrieval package is vendored
+# into the Space at push time (see hf/push_all.py), so it is importable without a pip install.
+gradio>=4.44
+spaces
+torch
+torchaudio
+numpy
+muq
+transformers
+hydra-core
+omegaconf
+huggingface_hub
+einops
+soundfile
+scikit-learn

steerable_retrieval/__init__.py

@@ -0,0 +1,8 @@
+"""Sparse steerable retrieval.
+
+Open-vocabulary concept control for dense music retrieval via sparse inversion in a
+trained sparse autoencoder. See :mod:`steerable_retrieval.steer` for the public API.
+"""
+
+__version__ = "0.1.0"
+

steerable_retrieval/assets/__init__.py

@@ -0,0 +1 @@
+"""Packaged data assets (e.g. the default Mahalanobis manifold prior)."""

steerable_retrieval/assets/muq_mulan_music4all_prior.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b37230f32cbe75f380d91a7d1226a2ec59ff87bde0a8c68c9a7c3787813adb11
+size 970844

steerable_retrieval/callbacks/__init__.py

@@ -0,0 +1,19 @@
+"""Callbacks for training and evaluation."""
+
+from steerable_retrieval.callbacks.energy import (
+    EnergyCallback,
+    ModalityScoreCallback,
+    SimLoggerCallback,
+)
+from steerable_retrieval.callbacks.save import SaveActivationsCallback
+from steerable_retrieval.callbacks.alignment import CKNNACallback
+from steerable_retrieval.callbacks.utils import BaseCallback
+
+__all__ = [
+    "EnergyCallback",
+    "ModalityScoreCallback",
+    "SimLoggerCallback",
+    "SaveActivationsCallback",
+    "CKNNACallback",
+    "BaseCallback",
+]

steerable_retrieval/callbacks/alignment.py

@@ -0,0 +1,253 @@
+"""
+CKNNACallback — Centered Kernel Nearest-Neighbor Alignment between
+SAE activations and input embeddings.
+
+Logs CKNNA scores for four cross-space pairs:
+  - text activations  <-> text embeddings
+  - audio activations <-> audio embeddings
+  - text activations  <-> audio embeddings
+  - audio activations <-> text embeddings
+
+Reads from pl_module.{val,test}_{activations,embeddings} populated by
+LightningSAE._eval_step.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Dict, List, Optional
+
+import torch
+from lightning.pytorch import Trainer
+from lightning.pytorch.core import LightningModule
+
+from steerable_retrieval.callbacks.utils import BaseCallback
+from steerable_retrieval.callbacks.energy import _get_dataset_name
+
+log = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# HSIC helpers (self-contained, no external deps)
+# ---------------------------------------------------------------------------
+
+def _hsic_unbiased(K: torch.Tensor, L: torch.Tensor) -> torch.Tensor:
+    """Unbiased HSIC estimator (Song et al., JMLR 2012, Eq. 5)."""
+    m = K.shape[0]
+    K_tilde = K.clone().fill_diagonal_(0)
+    L_tilde = L.clone().fill_diagonal_(0)
+    hsic = (
+        (K_tilde * L_tilde.T).sum()
+        + K_tilde.sum() * L_tilde.sum() / ((m - 1) * (m - 2))
+        - 2 * (K_tilde @ L_tilde).sum() / (m - 2)
+    )
+    return hsic / (m * (m - 3))
+
+
+def _hsic_biased(K: torch.Tensor, L: torch.Tensor) -> torch.Tensor:
+    """Biased HSIC (original CKA)."""
+    H = torch.eye(K.shape[0], dtype=K.dtype, device=K.device) - 1.0 / K.shape[0]
+    return torch.trace(K @ H @ L @ H)
+
+
+# ---------------------------------------------------------------------------
+# CKNNA metric
+# ---------------------------------------------------------------------------
+
+def cknna(
+    feats_A: torch.Tensor,
+    feats_B: torch.Tensor,
+    topk: int,
+    distance_agnostic: bool = False,
+    unbiased: bool = True,
+) -> float:
+    """
+    Centered Kernel Nearest-Neighbor Alignment.
+
+    Args:
+        feats_A: [N, D_A] feature matrix.
+        feats_B: [N, D_B] feature matrix.
+        topk:    number of nearest neighbors (>= 2).
+        distance_agnostic: if True, use binary neighbor overlap only.
+        unbiased: if True, use unbiased HSIC and exclude self-similarities.
+
+    Returns:
+        CKNNA score (float).
+    """
+    n = feats_A.shape[0]
+    if topk < 2:
+        raise ValueError("CKNNA requires topk >= 2")
+    topk = min(topk, n - 1) if unbiased else min(topk, n)
+
+    K = feats_A @ feats_A.T
+    L = feats_B @ feats_B.T
+    device = feats_A.device
+
+    hsic_fn = _hsic_unbiased if unbiased else _hsic_biased
+
+    def similarity(K_: torch.Tensor, L_: torch.Tensor, k: int) -> torch.Tensor:
+        if unbiased:
+            K_hat = K_.clone().fill_diagonal_(float("-inf"))
+            L_hat = L_.clone().fill_diagonal_(float("-inf"))
+        else:
+            K_hat, L_hat = K_, L_
+
+        _, topk_K_idx = torch.topk(K_hat, k, dim=1)
+        _, topk_L_idx = torch.topk(L_hat, k, dim=1)
+
+        mask_K = torch.zeros(n, n, device=device).scatter_(1, topk_K_idx, 1.0)
+        mask_L = torch.zeros(n, n, device=device).scatter_(1, topk_L_idx, 1.0)
+        mask = mask_K * mask_L
+
+        if distance_agnostic:
+            return (mask.sum()).float()
+        return hsic_fn(mask * K_, mask * L_)
+
+    sim_kl = similarity(K, L, topk)
+    sim_kk = similarity(K, K, topk)
+    sim_ll = similarity(L, L, topk)
+
+    denom = (torch.sqrt(sim_kk * sim_ll) + 1e-6).item()
+    return sim_kl.item() / denom
+
+
+# ---------------------------------------------------------------------------
+# Callback
+# ---------------------------------------------------------------------------
+
+class CKNNACallback(BaseCallback):
+    """
+    Computes CKNNA between SAE activations and input embeddings at the end of
+    each validation / test epoch.
+
+    Pairs logged:
+      - text_act   <-> text_emb
+      - audio_act  <-> audio_emb
+      - text_act   <-> audio_emb
+      - audio_act  <-> text_emb
+    """
+
+    def __init__(
+        self,
+        enable_on_validation: bool = True,
+        enable_on_test: bool = True,
+        every_n_steps: int = None,
+        every_n_epochs: int = 1,
+        topk: int = 10,
+        max_samples: int = 2048,
+        distance_agnostic: bool = False,
+        unbiased: bool = True,
+        prefix: str = "CKNNA",
+    ):
+        super().__init__(every_n_steps=every_n_steps, every_n_epochs=every_n_epochs)
+        self.enable_on_validation = enable_on_validation
+        self.enable_on_test = enable_on_test
+        self.topk = topk
+        self.max_samples = max_samples
+        self.distance_agnostic = distance_agnostic
+        self.unbiased = unbiased
+        self.prefix = prefix
+
+    def on_validation_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
+        if not self.enable_on_validation:
+            return
+        if not (self._check_step(trainer, pl_module) or self._check_epoch(trainer, pl_module)):
+            return
+        self._compute_and_log(trainer, pl_module, mode="val")
+
+    def on_test_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
+        if not self.enable_on_test:
+            return
+        if not (self._check_step(trainer, pl_module) or self._check_epoch(trainer, pl_module)):
+            return
+        self._compute_and_log(trainer, pl_module, mode="test")
+
+    @staticmethod
+    @torch.no_grad()
+    def _compute(
+        Za: Optional[torch.Tensor],
+        Zt: Optional[torch.Tensor],
+        Ea: Optional[torch.Tensor],
+        Et: Optional[torch.Tensor],
+        topk: int = 10,
+        max_samples: int = 2048,
+        distance_agnostic: bool = False,
+        unbiased: bool = True,
+    ) -> Dict[str, float]:
+        """
+        Pure computation — no trainer / pl_module / logging.
+
+        Args:
+            Za: audio activations [N_a, C] or None
+            Zt: text activations  [N_t, C] or None
+            Ea: audio embeddings  [N_a, D] or None
+            Et: text embeddings   [N_t, D] or None
+            topk: number of nearest neighbors for CKNNA
+            max_samples: subsample limit
+            distance_agnostic: binary neighbor overlap only
+            unbiased: use unbiased HSIC
+
+        Returns:
+            dict mapping pair label -> CKNNA score (float).
+        """
+        pairs = [
+            (Za, Ea, "audio_act__audio_emb"),
+            (Zt, Et, "text_act__text_emb"),
+            (Zt, Ea, "text_act__audio_emb"),
+            (Za, Et, "audio_act__text_emb"),
+        ]
+
+        scores = {}
+        for feats_A, feats_B, label in pairs:
+            if feats_A is None or feats_B is None:
+                continue
+
+            n = min(feats_A.size(0), feats_B.size(0))
+            fA, fB = feats_A[:n], feats_B[:n]
+
+            if n > max_samples:
+                idx = torch.randperm(n, device=fA.device)[:max_samples]
+                fA, fB = fA[idx], fB[idx]
+
+            try:
+                scores[label] = cknna(
+                    fA.float(), fB.float(),
+                    topk=topk,
+                    distance_agnostic=distance_agnostic,
+                    unbiased=unbiased,
+                )
+            except Exception as e:
+                log.warning(f"[CKNNACallback] Failed for {label}: {e}")
+
+        return scores
+
+    @torch.no_grad()
+    def _compute_and_log(self, trainer: Trainer, pl_module: LightningModule, mode: str):
+        """Gather data from module, call _compute, then log."""
+        from steerable_retrieval.callbacks.energy import _resolve_modality_tensors
+
+        all_acts = getattr(pl_module, f"{mode}_activations", {})
+        all_embs = getattr(pl_module, f"{mode}_embeddings", {})
+        device = pl_module.device
+
+        for dataloader_idx in all_acts:
+            acts_dl = all_acts.get(dataloader_idx, {})
+            embs_dl = all_embs.get(dataloader_idx, {})
+
+            Za = _resolve_modality_tensors(acts_dl, "audio", device, trainer)
+            Zt = _resolve_modality_tensors(acts_dl, "text", device, trainer)
+            Ea = _resolve_modality_tensors(embs_dl, "audio", device, trainer)
+            Et = _resolve_modality_tensors(embs_dl, "text", device, trainer)
+
+            scores = self._compute(
+                Za, Zt, Ea, Et,
+                topk=self.topk,
+                max_samples=self.max_samples,
+                distance_agnostic=self.distance_agnostic,
+                unbiased=self.unbiased,
+            )
+
+            dataset_name = _get_dataset_name(trainer, dataloader_idx, mode)
+            log_prefix = f"{self.prefix}/{dataset_name}"
+            for label, score in scores.items():
+                pl_module.log(f"{log_prefix}/{label}", score, prog_bar=False, sync_dist=True)

steerable_retrieval/callbacks/energy.py

@@ -0,0 +1,1543 @@
+"""
+Callbacks implementing Energy / Modality / Bridge / Cross-run Stability metrics
+for LightningSAE (spamr) based on the PCA/USAE paper (2504.11695v4).
+"""
+
+from __future__ import annotations
+
+import copy
+import logging
+import os
+from typing import Dict, List, Optional
+
+import torch
+import torch.nn.functional as F
+from lightning.pytorch import Trainer
+from lightning.pytorch.core import LightningModule
+
+from steerable_retrieval.callbacks.utils import BaseCallback
+
+import plotly.graph_objects as go
+
+log = logging.getLogger(__name__)
+
+# Bar plot styling: normalized neuron index (0-1), fixed width, black outline
+BAR_WIDTH = 0.005
+BAR_MARKER_LINE = dict(width=1, color="black")
+
+
+def _normalized_neuron_x(n: int):
+    """Return x-axis values in [0, 1] for n neurons."""
+    if n <= 1:
+        return [0.5] if n == 1 else []
+    return [i / (n - 1) for i in range(n)]
+
+
+def _get_wandb_logger(trainer: Trainer):
+    """Get wandb logger from trainer if available."""
+    try:
+        from lightning.pytorch.loggers import WandbLogger
+        for logger in trainer.loggers:
+            if isinstance(logger, WandbLogger):
+                return logger
+    except ImportError:
+        pass
+    return None
+
+
+def _get_dataset_name(trainer: Trainer, dataloader_idx: int, mode: str) -> str:
+    """Get dataset name for a dataloader from datamodule.names, or fallback to dataloader_{idx}."""
+    dm = getattr(trainer, "datamodule", None)
+    if dm is not None and hasattr(dm, "names"):
+        names = dm.names
+        if isinstance(names, dict) and mode in names:
+            name = names[mode].get(dataloader_idx)
+            if name is not None:
+                return str(name)
+    return f"dataloader_{dataloader_idx}"
+
+
+def _log_plotly_to_wandb(trainer: Trainer, key: str, fig):
+    """Log a plotly figure to wandb if available."""
+    wandb_logger = _get_wandb_logger(trainer)
+    if wandb_logger is None:
+        return
+    try:
+        import wandb
+        if wandb.run is not None:
+            wandb.log({key: fig}, step=trainer.global_step)
+    except ImportError:
+        pass
+
+
+def gather_tensor_if_distributed(tensor: torch.Tensor, trainer: Trainer) -> torch.Tensor:
+    """
+    Gather tensor from all processes if distributed training is enabled.
+    Uses PyTorch Lightning's strategy for distributed operations.
+    """
+    if trainer.world_size <= 1:
+        return tensor
+
+    try:
+        if hasattr(trainer.strategy, "all_gather"):
+            gathered = trainer.strategy.all_gather(tensor, sync_grads=False)
+            if isinstance(gathered, (list, tuple)):
+                gathered = torch.cat(gathered, dim=0)
+            elif isinstance(gathered, torch.Tensor):
+                if gathered.dim() > tensor.dim():
+                    gathered = gathered.view(-1, *gathered.shape[2:])
+            return gathered
+    except (AttributeError, NotImplementedError):
+        pass
+
+    import torch.distributed as dist
+    if not dist.is_initialized():
+        return tensor
+
+    device = tensor.device
+
+    local_size = torch.tensor([tensor.shape[0]], device=device, dtype=torch.long)
+    sizes = [torch.zeros_like(local_size) for _ in range(trainer.world_size)]
+    dist.all_gather(sizes, local_size)
+    sizes = [s.item() for s in sizes]
+    max_size = max(sizes)
+
+    if tensor.shape[0] < max_size:
+        padding_shape = list(tensor.shape)
+        padding_shape[0] = max_size - tensor.shape[0]
+        padding = torch.zeros(padding_shape, device=device, dtype=tensor.dtype)
+        tensor = torch.cat([tensor, padding], dim=0)
+
+    gathered_tensors = [torch.zeros_like(tensor) for _ in range(trainer.world_size)]
+    dist.all_gather(gathered_tensors, tensor)
+
+    gathered_list = []
+    for i, gathered_tensor in enumerate(gathered_tensors):
+        gathered_list.append(gathered_tensor[:sizes[i]])
+
+    return torch.cat(gathered_list, dim=0)
+
+
+def get_dictionary_from_lightningsae(pl_module: LightningModule) -> torch.Tensor:
+    """
+    Returns dictionary D as [C, d] where each row is a concept atom in embedding space.
+
+    The SAE decoder has W_dec of shape [dict_size, act_size] = [C, d].
+    """
+    if not hasattr(pl_module, "sae_decoder"):
+        raise AttributeError("Expected pl_module.sae_decoder to exist (LightningSAE).")
+
+    W = getattr(pl_module.sae_decoder, "W_dec", None)
+    if W is None or not torch.is_tensor(W):
+        raise AttributeError("Expected pl_module.sae_decoder.W_dec to be a Tensor/Parameter.")
+
+    return W.contiguous()
+
+
+def load_state_dict_any(path: str, map_location="cpu") -> Dict[str, torch.Tensor]:
+    """
+    With torch_s3_connector installed, torch.load("s3://...") should work.
+
+    Supports:
+      - raw state_dict
+      - Lightning checkpoint dict containing 'state_dict'
+    """
+    if str(path).startswith("s3://"):
+        from urllib.parse import urlparse
+
+        import boto3
+
+        parsed = urlparse(str(path))
+        bucket = parsed.netloc
+        key = parsed.path.lstrip("/")
+        if not bucket or not key:
+            raise ValueError(f"Invalid S3 checkpoint URI: {path}")
+
+        cache_root = os.environ.get("SPAMR_CHECKPOINT_CACHE_DIR", os.path.expanduser("~/.cache/steerable_retrieval/checkpoints"))
+        cache_path = os.path.join(cache_root, bucket, key)
+        force_refresh = str(os.environ.get("SPAMR_CHECKPOINT_CACHE_REFRESH", "0")).lower() in {"1", "true", "yes"}
+        os.makedirs(os.path.dirname(cache_path), exist_ok=True)
+
+        if force_refresh or not os.path.exists(cache_path):
+            boto3.client("s3").download_file(bucket, key, cache_path)
+            log.info(f"Downloaded checkpoint to local cache: {cache_path}")
+        else:
+            log.info(f"Using cached checkpoint: {cache_path}")
+
+        # Checkpoints in this project may include OmegaConf objects in metadata.
+        # PyTorch>=2.6 defaults to weights_only=True, which rejects those objects.
+        obj = torch.load(cache_path, map_location=map_location, weights_only=False)
+    else:
+        obj = torch.load(path, map_location=map_location, weights_only=False)
+    if isinstance(obj, dict) and "state_dict" in obj and isinstance(obj["state_dict"], dict):
+        return obj["state_dict"]
+    if isinstance(obj, dict):
+        return obj
+    raise ValueError(f"Unsupported checkpoint object at {path}: {type(obj)}")
+
+
+def _safe_name(path: str) -> str:
+    base = path.rstrip("/").split("/")[-1]
+    base = base.replace(".", "_").replace("-", "_").replace("=", "_")
+    if len(base) > 48:
+        base = base[:48]
+    return base
+
+
+def _cat_or_none(xs: List[torch.Tensor], device: torch.device) -> Optional[torch.Tensor]:
+    if not xs:
+        return None
+    return torch.cat([x.to(device) for x in xs], dim=0)
+
+
+def _resolve_modality_tensors(
+    store: Dict,
+    modality: str,
+    device: torch.device,
+    trainer: Optional["Trainer"] = None,
+) -> Optional[torch.Tensor]:
+    """
+    Resolve a modality's data from a per-dataloader storage dict.
+    Handles both list-of-tensors and already-concatenated tensor formats.
+    Optionally gathers across ranks when trainer is provided.
+    """
+    data = store.get(modality)
+    if data is None:
+        return None
+    if isinstance(data, list):
+        t = _cat_or_none(data, device=device)
+    elif isinstance(data, torch.Tensor):
+        t = data.to(device)
+    else:
+        return None
+    if t is not None and trainer is not None:
+        t = gather_tensor_if_distributed(t, trainer)
+    return t
+
+
+# -----------------------------------------------------------------------------
+# EnergyCallback
+# -----------------------------------------------------------------------------
+class EnergyCallback(BaseCallback):
+    """
+    Energy_i = E[z_i] computed from logged SAE activations.
+
+    Logs:
+      - {mode}/energy_sum
+      - {mode}/energy_mean
+      - {mode}/energy_topK_frac   (fraction of total energy contained in top-K concepts)
+
+    Also stores:
+      pl_module._last_energy = {'combined': E, 'audio': Ea or None, 'text': Et or None}
+    """
+
+    def __init__(
+        self,
+        enable_on_validation: bool = True,
+        enable_on_test: bool = True,
+        every_n_steps: int = None,
+        every_n_epochs: int = 1,
+        topk: int = 512,
+        prefix: str = "energy",
+    ):
+        super().__init__(every_n_steps=every_n_steps, every_n_epochs=every_n_epochs)
+        self.enable_on_validation = enable_on_validation
+        self.enable_on_test = enable_on_test
+        self.topk = topk
+        self.prefix = prefix
+
+    def on_validation_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
+        if not self.enable_on_validation:
+            return
+        if not (self._check_step(trainer, pl_module) or self._check_epoch(trainer, pl_module)):
+            return
+        self._compute_and_log(trainer, pl_module, mode="val")
+
+    def on_test_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
+        if not self.enable_on_test:
+            return
+        self._compute_and_log(trainer, pl_module, mode="test")
+
+    @staticmethod
+    @torch.no_grad()
+    def _compute(
+        Za: Optional[torch.Tensor],
+        Zt: Optional[torch.Tensor],
+        topk: int = 512,
+    ) -> Optional[Dict]:
+        """
+        Pure computation — no trainer / pl_module / logging.
+
+        Args:
+            Za: audio activations [N_a, C] or None
+            Zt: text activations  [N_t, C] or None
+            topk: number of top concepts for energy fraction
+
+        Returns:
+            dict with keys: E, Ea, Et, total, frac_topk, k, Za, Zt
+            or None if both inputs are None.
+        """
+        if Za is None and Zt is None:
+            return None
+
+        Ea = Za.mean(dim=0) if Za is not None else None
+        Et = Zt.mean(dim=0) if Zt is not None else None
+
+        if Ea is not None and Et is not None:
+            E = 0.5 * (Ea + Et)
+        else:
+            E = Ea if Ea is not None else Et
+
+        total = E.sum().clamp_min(1e-12)
+        k = min(topk, E.numel())
+        frac_topk = (torch.topk(E, k=k).values.sum() / total).item()
+
+        return dict(E=E, Ea=Ea, Et=Et, total=total, frac_topk=frac_topk, k=k, Za=Za, Zt=Zt)
+
+    @torch.no_grad()
+    def _compute_and_log(self, trainer: Trainer, pl_module: LightningModule, mode: str):
+        """Gather data from module, call _compute, then log."""
+        all_acts = pl_module.val_activations if mode == "val" else pl_module.test_activations
+        device = pl_module.device
+
+        for dataloader_idx, acts in all_acts.items():
+            Za = _resolve_modality_tensors(acts, "audio", device, trainer)
+            Zt = _resolve_modality_tensors(acts, "text", device, trainer)
+
+            result = self._compute(Za, Zt, topk=self.topk)
+            if result is None:
+                log.warning(f"[EnergyCallback] No activations found for mode={mode}, dataloader_idx={dataloader_idx}.")
+                continue
+
+            E, Ea, Et = result["E"], result["Ea"], result["Et"]
+            dataset_name = _get_dataset_name(trainer, dataloader_idx, mode)
+            log_prefix = f"Energy/{dataset_name}"
+            pl_module.log(f"{log_prefix}/sum", result["total"].item(), sync_dist=True)
+            pl_module.log(f"{log_prefix}/mean", E.mean().item(), sync_dist=True)
+            pl_module.log(f"{log_prefix}/top{result['k']}_frac", result["frac_topk"], prog_bar=True, sync_dist=True)
+
+            if dataloader_idx == 0 or len(all_acts) == 1:
+                pl_module._last_energy = {
+                    "combined": E.detach(),
+                    "audio": Ea.detach() if Ea is not None else None,
+                    "text": Et.detach() if Et is not None else None,
+                }
+
+            if Za is not None and Zt is not None:
+                self._log_energy_distribution_combined(trainer, Za, Zt, mode, dataset_name)
+            for activations, name in zip([Za, Zt], ["audio", "text"]):
+                if activations is not None:
+                    self._log_energy_distribution(trainer, activations, name, mode, dataset_name)
+
+            self._log_cumulative_energy_plot(trainer, pl_module, E, mode, dataset_name)
+
+    def _log_energy_distribution_combined(
+        self,
+        trainer: Trainer,
+        Za: torch.Tensor,
+        Zt: torch.Tensor,
+        mode: str,
+        dataset_name: str,
+    ):
+        """Stacked bar: mean activation per neuron from audio (blue) and text (red)."""
+        try:
+            import plotly.graph_objects as go
+        except ImportError:
+            log.debug("plotly not available, skipping energy distribution combined plot")
+            return
+
+        mean_a = Za.detach().cpu().mean(dim=0)
+        mean_t = Zt.detach().cpu().mean(dim=0)
+        n_concepts = mean_a.shape[0]
+        x_norm = _normalized_neuron_x(n_concepts)
+
+        fig = go.Figure()
+        fig.add_trace(go.Bar(
+            x=x_norm,
+            y=mean_a.numpy().tolist(),
+            name="audio",
+            width=BAR_WIDTH,
+            marker=dict(color="blue", line=BAR_MARKER_LINE),
+            hovertemplate="Neuron: %{x:.4f}<br>Audio: %{y:.4f}<extra></extra>",
+        ))
+        fig.add_trace(go.Bar(
+            x=x_norm,
+            y=mean_t.numpy().tolist(),
+            name="text",
+            width=BAR_WIDTH,
+            marker=dict(color="red", line=BAR_MARKER_LINE),
+            hovertemplate="Neuron: %{x:.4f}<br>Text: %{y:.4f}<extra></extra>",
+        ))
+        fig.update_layout(
+            barmode="stack",
+            title=f"Energy Distribution — combined ({mode})",
+            xaxis_title="Neuron Index (normalized 0–1)",
+            yaxis_title="Mean Activation",
+            template="plotly_white",
+        )
+        _log_plotly_to_wandb(trainer, f"Energy Distribution/{dataset_name}/energy_distribution_combined", fig)
+
+    def _log_energy_distribution(
+        self,
+        trainer: Trainer,
+        activations: torch.Tensor,
+        name: str,
+        mode: str,
+        dataset_name: str,
+    ):
+        """Plot average activation per neuron (bar height = mean, error bars = std from raw data)."""
+        try:
+            import plotly.graph_objects as go
+        except ImportError:
+            log.debug("plotly not available, skipping energy distribution plot")
+            return
+
+        acts = activations.detach().cpu()
+        mean_per_neuron = acts.mean(dim=0)
+        std_per_neuron = acts.std(dim=0)
+        n_concepts = mean_per_neuron.shape[0]
+
+        logging.info(f"dataset_name: {dataset_name}, modality: {name}, mean_per_neuron: {mean_per_neuron.shape}, std_per_neuron: {std_per_neuron.shape}, n_concepts: {n_concepts}")
+        
+        # check if nans in activations
+
+        if torch.isnan(mean_per_neuron).any() or torch.isnan(std_per_neuron).any():
+            log.warning(f"NaNs found in activations for dataset_name: {dataset_name}, modality: {name}")
+            
+        logging.info(f"neurons with nans: {torch.isnan(mean_per_neuron).nonzero()}, {torch.isnan(std_per_neuron).nonzero()}")
+        logging.info(f'mean_per_neuron: {mean_per_neuron}, std_per_neuron: {std_per_neuron}')
+
+
+        # Modality color: blue for audio, red for text
+        color = "blue" if name == "audio" else "red"
+        x_norm = _normalized_neuron_x(n_concepts)
+        fig = go.Figure()
+        fig.add_trace(go.Bar(
+            x=x_norm,
+            y=mean_per_neuron.numpy().tolist(),
+            name=name,
+            width=BAR_WIDTH,
+            marker=dict(color=color, line=BAR_MARKER_LINE),
+            hovertemplate="Neuron: %{x:.4f}<br>Mean: %{y:.4f}<extra></extra>",
+        ))
+        fig.update_layout(
+            title=f"Energy Distribution — {name} ({mode})",
+            xaxis_title="Neuron Index (normalized 0–1)",
+            yaxis_title="Mean Activation",
+            template="plotly_white",
+        )
+        _log_plotly_to_wandb(trainer, f"Energy Distribution/{dataset_name}/energy_distribution_{name}", fig)
+
+    def _log_cumulative_energy_plot(
+        self,
+        trainer: Trainer,
+        pl_module: LightningModule,
+        E: torch.Tensor,
+        mode: str,
+        dataset_name: str,
+        log_scale: bool = True,
+    ):
+        """
+        Log cumulative energy plot: concepts sorted by energy (descending),
+        x-axis = concept rank, y-axis = cumulative normalized energy.
+        If log_scale is True (default), x-axis is log-scaled.
+        """
+        try:
+            import plotly.graph_objects as go
+        except ImportError:
+            log.debug("plotly not available, skipping cumulative energy plot")
+            return
+
+        # Sort energies descending
+        E_sorted, _ = torch.sort(E, descending=True)
+        E_sorted = E_sorted.detach().cpu()
+        total_energy = E_sorted.sum().clamp_min(1e-12)
+
+        # Compute cumulative normalized energy
+        cumulative = torch.cumsum(E_sorted, dim=0) / total_energy
+        cumulative = cumulative.numpy()
+
+        n_concepts = len(cumulative)
+        x = list(range(1, n_concepts + 1))
+
+        fig = go.Figure()
+        fig.add_trace(go.Scatter(
+            x=x,
+            y=cumulative.tolist(),
+            mode="lines",
+            name="Cumulative Energy",
+            line=dict(color="blue", width=2),
+        ))
+        fig.update_layout(
+            title=f"Cumulative Normalized Energy ({mode})",
+            xaxis_title="Number of Concepts (ranked by energy)",
+            yaxis_title="Cumulative Energy Fraction",
+            yaxis_range=[0, 1.05],
+            template="plotly_white",
+        )
+
+        if log_scale:
+            fig.update_xaxes(type="log")
+
+        plot_key = f"Energy Distribution/{dataset_name}/cumulative_plot"
+        _log_plotly_to_wandb(trainer, plot_key, fig)
+
+
+# -----------------------------------------------------------------------------
+# ModalityScoreCallback
+# -----------------------------------------------------------------------------
+class ModalityScoreCallback(BaseCallback):
+    """
+    ModalityScore_i = E_audio[z_i] / (E_audio[z_i] + E_text[z_i])
+
+    Logs summary stats:
+      - mean modality score
+      - fraction mostly-audio (>0.9)
+      - fraction mostly-text (<0.1)
+
+    Stores:
+      pl_module._last_modality_score = score
+    """
+
+    def __init__(
+        self,
+        enable_on_validation: bool = True,
+        enable_on_test: bool = True,
+        every_n_steps: int = None,
+        every_n_epochs: int = 1,
+        eps: float = 1e-5,
+        prefix: str = "modality",
+    ):
+        super().__init__(every_n_steps=every_n_steps, every_n_epochs=every_n_epochs)
+        self.enable_on_validation = enable_on_validation
+        self.enable_on_test = enable_on_test
+        self.eps = eps
+        self.prefix = prefix
+
+    def on_validation_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
+        if not self.enable_on_validation:
+            return
+        if not (self._check_step(trainer, pl_module) or self._check_epoch(trainer, pl_module)):
+            return
+        self._compute_and_log(trainer, pl_module, mode="val")
+
+    def on_test_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
+        if not self.enable_on_test:
+            return
+        if not (self._check_step(trainer, pl_module) or self._check_epoch(trainer, pl_module)):
+            return
+        self._compute_and_log(trainer, pl_module, mode="test")
+
+    @staticmethod
+    @torch.no_grad()
+    def _compute(
+        Za: torch.Tensor,
+        Zt: torch.Tensor,
+        eps: float = 1e-5,
+    ) -> Dict:
+        """
+        Pure computation — no trainer / pl_module / logging.
+
+        Args:
+            Za: audio activations [N_a, C]
+            Zt: text activations  [N_t, C]
+            eps: epsilon for numerical stability
+
+        Returns:
+            dict with keys: score, E, Ea, Et, mean, frac_audio, frac_text
+        """
+        Ea = Za.mean(dim=0)
+        Et = Zt.mean(dim=0)
+        score = Ea / (Ea + Et + eps)
+        E = 0.5 * (Ea + Et)
+        return dict(
+            score=score,
+            E=E,
+            Ea=Ea,
+            Et=Et,
+            mean=score.mean().item(),
+            frac_audio=((score > 0.9).float().mean().item()),
+            frac_text=((score < 0.1).float().mean().item()),
+        )
+
+    @torch.no_grad()
+    def _compute_and_log(self, trainer: Trainer, pl_module: LightningModule, mode: str):
+        """Gather data from module, call _compute, then log."""
+        all_acts = pl_module.val_activations if mode == "val" else pl_module.test_activations
+        device = pl_module.device
+
+        for dataloader_idx, acts in all_acts.items():
+            Za = _resolve_modality_tensors(acts, "audio", device, trainer)
+            Zt = _resolve_modality_tensors(acts, "text", device, trainer)
+
+            if Za is None or Zt is None:
+                log.warning(f"[ModalityScoreCallback] Missing modality activations for mode={mode}, dataloader_idx={dataloader_idx}.")
+                continue
+
+            result = self._compute(Za, Zt, eps=self.eps)
+            score, E = result["score"], result["E"]
+
+            dataset_name = _get_dataset_name(trainer, dataloader_idx, mode)
+            log_prefix = f"Modality Score/{dataset_name}"
+            pl_module.log(f"{log_prefix}/mean", result["mean"], prog_bar=True, sync_dist=True)
+            pl_module.log(f"{log_prefix}/frac_audio_gt0.9", result["frac_audio"], sync_dist=True)
+            pl_module.log(f"{log_prefix}/frac_text_lt0.1", result["frac_text"], sync_dist=True)
+
+            if dataloader_idx == 0 or len(all_acts) == 1:
+                pl_module._last_modality_score = score.detach()
+
+            self._log_modality_histograms(trainer, score, E, mode, dataset_name)
+            self._log_modality_score_distribution(trainer, score, mode, dataset_name)
+
+    @torch.no_grad()
+    def _log_modality_score_distribution(
+        self,
+        trainer: Trainer,
+        score: torch.Tensor,
+        mode: str,
+        dataset_name: str,
+    ):
+        """Plot modality score per neuron index using plotly."""
+        try:
+            import plotly.graph_objects as go
+        except ImportError:
+            log.debug("plotly not available, skipping modality score distribution plot")
+            return
+
+        score = score.detach().cpu()
+        n = score.shape[0]
+        x_norm = _normalized_neuron_x(n)
+        fig = go.Figure()
+        fig.add_trace(go.Bar(
+            x=x_norm,
+            y=score.numpy().tolist(),
+            name="Modality Score",
+            width=BAR_WIDTH,
+            marker=dict(color="steelblue", line=BAR_MARKER_LINE),
+        ))
+        fig.update_layout(
+            title=f"Modality Score Distribution ({mode})",
+            xaxis_title="Neuron Index (normalized 0–1)",
+            yaxis_title="Modality Score",
+            template="plotly_white",
+        )
+        _log_plotly_to_wandb(trainer, f"Modality Score/{dataset_name}/modality_score_distribution", fig)
+
+    def _log_modality_histograms(
+        self,
+        trainer: Trainer,
+        score: torch.Tensor,
+        E: torch.Tensor,
+        mode: str,
+        dataset_name: str,
+    ):
+        """
+        Log two histograms of modality scores per concept:
+        - Unweighted: count of concepts per bin (normalized to proportion).
+        - Energy-weighted: energy in each bin (normalized by total energy).
+        """
+        try:
+            import numpy as np
+            import plotly.graph_objects as go
+        except ImportError:
+            log.debug("numpy/plotly not available, skipping modality histograms")
+            return
+
+        score_np = score.detach().cpu().float().numpy()
+        E_np = E.detach().cpu().float().numpy()
+        bins = np.linspace(0.0, 1.0, 51)
+        bin_centers = 0.5 * (bins[:-1] + bins[1:])
+
+        # Unweighted: count per bin, normalize to proportion
+        counts, _ = np.histogram(score_np, bins=bins)
+        total = max(counts.sum(), 1e-12)
+        density_unweighted = (counts / total).tolist()
+
+        # Energy-weighted: sum of energy per bin, normalize by total energy
+        counts_w, _ = np.histogram(score_np, bins=bins, weights=E_np)
+        total_E = max(E_np.sum(), 1e-12)
+        density_weighted = (counts_w / total_E).tolist()
+
+        layout_base = dict(
+            xaxis_title="Modality Score (0=text, 1=audio)",
+            yaxis_title="Proportion",
+            template="plotly_white",
+            bargap=0.1,
+        )
+
+        fig_unweighted = go.Figure()
+        fig_unweighted.add_trace(go.Bar(
+            x=bin_centers.tolist(),
+            y=density_unweighted,
+            marker=dict(color="steelblue", line=BAR_MARKER_LINE),
+            width=0.005,
+        ))
+        fig_unweighted.update_layout(
+            title=f"Modality Score Distribution — Unweighted ({mode})",
+            **layout_base,
+        )
+        fig_unweighted.update_yaxes(rangemode="tozero")
+        _log_plotly_to_wandb(
+            trainer,
+            f"Modality Score/{dataset_name}/distribution_unweighted",
+            fig_unweighted,
+        )
+
+        fig_weighted = go.Figure()
+        fig_weighted.add_trace(go.Bar(
+            x=bin_centers.tolist(),
+            y=density_weighted,
+            marker=dict(color="coral", line=BAR_MARKER_LINE),
+            width=0.005,
+        ))
+        fig_weighted.update_layout(
+            title=f"Modality Score Distribution — Energy-Weighted ({mode})",
+            **layout_base,
+        )
+        fig_weighted.update_yaxes(rangemode="tozero")
+        _log_plotly_to_wandb(
+            trainer,
+            f"Modality Score/{dataset_name}/distribution_energy_weighted",
+            fig_weighted,
+        )
+
+
+# -----------------------------------------------------------------------------
+# BridgeScoreCallback
+# -----------------------------------------------------------------------------
+class BridgeScoreCallback(BaseCallback):
+    """
+    Bridge matrix:
+      B = E[z_a^T z_t] ⊙ (D D^T)
+    where:
+      - E[z_a^T z_t] is estimated as (Za^T Zt)/N for paired samples
+      - D D^T uses cosine similarity between dictionary atoms (rows of D)
+
+    Logs only matrix summaries (avoid logging full B):
+      - abs mean
+      - abs max
+      - abs topK mean
+
+    Stores:
+      pl_module._last_bridge = B
+    """
+
+    def __init__(
+        self,
+        enable_on_validation: bool = True,
+        enable_on_test: bool = True,
+        every_n_steps: int = None,
+        every_n_epochs: int = 1,
+        prefix: str = "bridge",
+        topk_edges: int = 200,
+    ):
+        super().__init__(every_n_steps=every_n_steps, every_n_epochs=every_n_epochs)
+        self.enable_on_validation = enable_on_validation
+        self.enable_on_test = enable_on_test
+        self.prefix = prefix
+        self.topk_edges = topk_edges
+
+    def on_validation_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
+        if not self.enable_on_validation:
+            return
+        if not (self._check_step(trainer, pl_module) or self._check_epoch(trainer, pl_module)):
+            return
+        self._compute_and_log(trainer, pl_module, mode="val")
+
+    def on_test_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
+        if not self.enable_on_test:
+            return
+        if not (self._check_step(trainer, pl_module) or self._check_epoch(trainer, pl_module)):
+            return
+        self._compute_and_log(trainer, pl_module, mode="test")
+
+    @staticmethod
+    @torch.no_grad()
+    def _compute(
+        Za: torch.Tensor,
+        Zt: torch.Tensor,
+        W_dec: torch.Tensor,
+        topk_edges: int = 200,
+    ) -> Dict:
+        """
+        Pure computation — no trainer / pl_module / logging.
+
+        Args:
+            Za: audio activations [N, C]
+            Zt: text activations  [N, C]  (must be same N)
+            W_dec: decoder dictionary [C, d]
+            topk_edges: number of top bridge edges
+
+        Returns:
+            dict with keys: B, abs_mean, abs_max, abs_topk_mean, k
+        """
+        if Za.shape[0] != Zt.shape[0]:
+            raise ValueError(f"BridgeScore requires paired N. Got Za={Za.shape}, Zt={Zt.shape}")
+
+        N = Za.shape[0]
+        coact = (Za.t() @ Zt) / max(N, 1)
+
+        Dn = F.normalize(W_dec, dim=-1)
+        align = Dn @ Dn.t()
+        B = coact * align
+
+        absB = B.abs()
+        k = min(topk_edges, absB.numel())
+        top_vals = torch.topk(absB.cpu().flatten(), k=k).values
+
+        return dict(
+            align=align,
+            coact=coact,
+            B=B,
+            abs_mean=absB.mean().item(),
+            abs_max=absB.max().item(),
+            abs_topk_mean=top_vals.mean().item(),
+            k=k,
+        )
+
+    @torch.no_grad()
+    def _compute_and_log(self, trainer: Trainer, pl_module: LightningModule, mode: str):
+        """Gather data from module, call _compute, then log."""
+        all_acts = pl_module.val_activations if mode == "val" else pl_module.test_activations
+        device = pl_module.device
+
+        bridges = {}
+
+        for dataloader_idx, acts in all_acts.items():
+            Za = _resolve_modality_tensors(acts, "audio", device, trainer)
+            Zt = _resolve_modality_tensors(acts, "text", device, trainer)
+
+            if Za is None or Zt is None:
+                log.warning(f"[BridgeScoreCallback] Missing modality activations for mode={mode}, dataloader_idx={dataloader_idx}.")
+                continue
+
+            W_dec = get_dictionary_from_lightningsae(pl_module).to(device)
+            result = self._compute(Za, Zt, W_dec, topk_edges=self.topk_edges)
+
+            dataset_name = _get_dataset_name(trainer, dataloader_idx, mode)
+            bridges[dataset_name] = {
+                "align": result["align"],
+                "coact": result["coact"],
+                "B": result["B"],
+            }
+            log_prefix = f"Bridge/{dataset_name}"
+            pl_module.log(f"{log_prefix}/abs_mean", result["abs_mean"], prog_bar=True, sync_dist=True)
+            pl_module.log(f"{log_prefix}/abs_max", result["abs_max"], sync_dist=True)
+            pl_module.log(f"{log_prefix}/abs_top{result['k']}_mean", result["abs_topk_mean"], sync_dist=True)
+
+            if dataloader_idx == 0 or len(all_acts) == 1:
+                pl_module._last_bridge = result["B"].detach()
+
+        # Preferred structure: per-dataset bridge tensors.
+        # Example: pl_module.bridges[dataset_name]["align" | "coact" | "B"]
+        pl_module.bridges = bridges
+
+
+# -----------------------------------------------------------------------------
+# ModalityClassifierCallback
+# -----------------------------------------------------------------------------
+class ModalityClassifierCallback(BaseCallback):
+    """
+    Per-concept modality classification accuracy.
+
+    For each concept, trains a logistic regression classifier to distinguish
+    audio vs text embeddings based on that concept's activation value.
+
+    Logs:
+      - Mean accuracy across all concepts
+      - Histogram of per-concept accuracies (unweighted)
+      - Histogram of per-concept accuracies weighted by energy
+
+    This helps identify which concepts are most discriminative for modality.
+    """
+
+    def __init__(
+        self,
+        enable_on_validation: bool = True,
+        enable_on_test: bool = True,
+        every_n_steps: int = None,
+        every_n_epochs: int = 1,
+        prefix: str = "modality_classifier",
+        max_samples: int = 10000,
+        test_fraction: float = 0.2,
+    ):
+        super().__init__(every_n_steps=every_n_steps, every_n_epochs=every_n_epochs)
+        self.enable_on_validation = enable_on_validation
+        self.enable_on_test = enable_on_test
+        self.prefix = prefix
+        self.max_samples = max_samples
+        self.test_fraction = test_fraction
+
+    def on_validation_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
+        if not self.enable_on_validation:
+            return
+        if not (self._check_step(trainer, pl_module) or self._check_epoch(trainer, pl_module)):
+            return
+        self._compute_and_log(trainer, pl_module, mode="val")
+
+    def on_test_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
+        if not self.enable_on_test:
+            return
+        if not (self._check_step(trainer, pl_module) or self._check_epoch(trainer, pl_module)):
+            return
+        self._compute_and_log(trainer, pl_module, mode="test")
+
+    @staticmethod
+    @torch.no_grad()
+    def _compute(
+        Za: torch.Tensor,
+        Zt: torch.Tensor,
+        max_samples: int = 10000,
+        test_fraction: float = 0.2,
+    ) -> Optional[Dict]:
+        """
+        Pure computation — no trainer / pl_module / logging.
+
+        Args:
+            Za: audio activations [N_a, C]
+            Zt: text activations  [N_t, C]
+            max_samples: subsample limit per modality
+            test_fraction: held-out fraction for classifier
+
+        Returns:
+            dict with keys: accuracies (np.ndarray [C]), E_np (np.ndarray [C]),
+            mean_acc, energy_weighted_mean_acc
+        """
+        try:
+            import numpy as np
+            from sklearn.linear_model import LogisticRegression
+            from sklearn.model_selection import train_test_split
+        except ImportError:
+            log.warning("[ModalityClassifierCallback] sklearn not available, skipping.")
+            return None
+
+        if Za.shape[0] > max_samples:
+            Za = Za[torch.randperm(Za.shape[0])[:max_samples]]
+        if Zt.shape[0] > max_samples:
+            Zt = Zt[torch.randperm(Zt.shape[0])[:max_samples]]
+
+        Za_np = Za.cpu().float().numpy()
+        Zt_np = Zt.cpu().float().numpy()
+        n_concepts = Za_np.shape[1]
+
+        Ea = Za.mean(dim=0)
+        Et = Zt.mean(dim=0)
+        E = 0.5 * (Ea + Et)
+        E_np = E.cpu().float().numpy()
+
+        y_audio = np.ones(Za_np.shape[0], dtype=np.int32)
+        y_text = np.zeros(Zt_np.shape[0], dtype=np.int32)
+
+        accuracies = []
+        for c in range(n_concepts):
+            X_c = np.concatenate([Za_np[:, c : c + 1], Zt_np[:, c : c + 1]], axis=0)
+            y_c = np.concatenate([y_audio, y_text], axis=0)
+            if X_c.std() < 1e-9:
+                accuracies.append(0.5)
+                continue
+            try:
+                X_train, X_test, y_train, y_test = train_test_split(
+                    X_c, y_c, test_size=test_fraction, stratify=y_c, random_state=42
+                )
+                clf = LogisticRegression(max_iter=200, solver="lbfgs")
+                clf.fit(X_train, y_train)
+                accuracies.append(clf.score(X_test, y_test))
+            except Exception:
+                accuracies.append(0.5)
+
+        accuracies = np.array(accuracies)
+        mean_acc = float(accuracies.mean())
+        ew_mean_acc = float((accuracies * E_np).sum() / max(E_np.sum(), 1e-12))
+
+        return dict(
+            accuracies=accuracies,
+            E_np=E_np,
+            mean_acc=mean_acc,
+            energy_weighted_mean_acc=ew_mean_acc,
+        )
+
+    @torch.no_grad()
+    def _compute_and_log(self, trainer: Trainer, pl_module: LightningModule, mode: str):
+        """Gather data from module, call _compute, then log."""
+        all_acts = pl_module.val_activations if mode == "val" else pl_module.test_activations
+        device = pl_module.device
+
+        for dataloader_idx, acts in all_acts.items():
+            Za = _resolve_modality_tensors(acts, "audio", device, trainer)
+            Zt = _resolve_modality_tensors(acts, "text", device, trainer)
+
+            if Za is None or Zt is None:
+                log.warning(f"[ModalityClassifierCallback] Missing modality activations for mode={mode}, dataloader_idx={dataloader_idx}.")
+                continue
+
+            result = self._compute(Za, Zt, max_samples=self.max_samples, test_fraction=self.test_fraction)
+            if result is None:
+                continue
+
+            dataset_name = _get_dataset_name(trainer, dataloader_idx, mode)
+            log_prefix = f"Modality Classifier/{dataset_name}"
+            pl_module.log(f"{log_prefix}/mean_acc", result["mean_acc"], prog_bar=True, sync_dist=True)
+            pl_module.log(f"{log_prefix}/energy_weighted_mean_acc", result["energy_weighted_mean_acc"], sync_dist=True)
+
+            self._log_accuracy_histograms(trainer, result["accuracies"], result["E_np"], mode, dataset_name)
+            self._log_per_neuron_accuracy_barplot(trainer, result["accuracies"], result["E_np"], mode, dataset_name)
+
+    def _log_accuracy_histograms(
+        self,
+        trainer: Trainer,
+        accuracies,
+        E,
+        mode: str,
+        dataset_name: str,
+    ):
+        """
+        Log two histograms of per-concept classification accuracies:
+        - Unweighted: proportion of concepts in each accuracy bin.
+        - Energy-weighted: proportion of energy in each accuracy bin.
+        """
+        try:
+            import numpy as np
+            import plotly.graph_objects as go
+        except ImportError:
+            log.debug("numpy/plotly not available, skipping accuracy histograms")
+            return
+
+        bins = np.linspace(0.0, 1.0, 51)
+        bin_centers = 0.5 * (bins[:-1] + bins[1:])
+
+        # Unweighted
+        counts, _ = np.histogram(accuracies, bins=bins)
+        total = max(counts.sum(), 1e-12)
+        density_unweighted = (counts / total).tolist()
+
+        # Energy-weighted
+        counts_w, _ = np.histogram(accuracies, bins=bins, weights=E)
+        total_E = max(E.sum(), 1e-12)
+        density_weighted = (counts_w / total_E).tolist()
+
+        layout_base = dict(
+            xaxis_title="Per-Concept Classification Accuracy",
+            yaxis_title="Proportion",
+            template="plotly_white",
+            bargap=0.1,
+        )
+
+        fig_unweighted = go.Figure()
+        fig_unweighted.add_trace(
+            go.Bar(
+                x=bin_centers.tolist(),
+                y=density_unweighted,
+                marker=dict(color="teal", line=BAR_MARKER_LINE),
+                width=BAR_WIDTH,
+            )
+        )
+        fig_unweighted.update_layout(
+            title=f"Per-Concept Modality Classification Accuracy — Unweighted ({mode})",
+            **layout_base,
+        )
+        fig_unweighted.update_yaxes(rangemode="tozero")
+        _log_plotly_to_wandb(
+            trainer,
+            f"Modality Classifier/{dataset_name}/acc_histogram_unweighted",
+            fig_unweighted,
+        )
+
+        fig_weighted = go.Figure()
+        fig_weighted.add_trace(
+            go.Bar(
+                x=bin_centers.tolist(),
+                y=density_weighted,
+                marker=dict(color="orange", line=BAR_MARKER_LINE),
+                width=BAR_WIDTH,
+            )
+        )
+        fig_weighted.update_layout(
+            title=f"Per-Concept Modality Classification Accuracy — Energy-Weighted ({mode})",
+            **layout_base,
+        )
+        fig_weighted.update_yaxes(rangemode="tozero")
+        _log_plotly_to_wandb(
+            trainer,
+            f"Modality Classifier/{dataset_name}/acc_histogram_energy_weighted",
+            fig_weighted,
+        )
+
+    def _log_per_neuron_accuracy_barplot(
+        self,
+        trainer: Trainer,
+        accuracies,
+        E,
+        mode: str,
+        dataset_name: str,
+    ):
+        """
+        Log a barplot showing accuracy for each neuron/concept.
+        X-axis: neuron index, Y-axis: classification accuracy.
+        """
+        try:
+            import numpy as np
+            import plotly.graph_objects as go
+        except ImportError:
+            log.debug("numpy/plotly not available, skipping per-neuron accuracy barplot")
+            return
+
+        n_concepts = len(accuracies)
+        x_norm = _normalized_neuron_x(n_concepts)
+
+        # Sort by accuracy (descending) for better visualization
+        sorted_idx = np.argsort(accuracies)[::-1]
+        sorted_accuracies = accuracies[sorted_idx]
+        sorted_energies = E[sorted_idx]
+
+        # Color by energy (normalized)
+        E_norm = sorted_energies / max(sorted_energies.max(), 1e-12)
+        
+        fig = go.Figure()
+        fig.add_trace(
+            go.Bar(
+                x=x_norm,
+                y=sorted_accuracies.tolist(),
+                width=0.001,
+                marker=dict(
+                    color=E_norm.tolist(),
+                    colorscale="Viridis",
+                    colorbar=dict(title="Normalized<br>Energy"),
+                    line=BAR_MARKER_LINE,
+                ),
+                hovertemplate="Neuron: %{x:.4f}<br>Accuracy: %{y:.3f}<extra></extra>",
+            )
+        )
+        fig.update_layout(
+            title=f"Per-Neuron Classification Accuracy — sorted by accuracy ({mode})",
+            xaxis_title="Neuron Index (normalized 0–1, sorted by accuracy)",
+            yaxis_title="Classification Accuracy",
+            template="plotly_white",
+            yaxis=dict(range=[0.5, 1]),
+            height=500,
+        )
+        _log_plotly_to_wandb(
+            trainer,
+            f"Modality Classifier/{dataset_name}/per_neuron_accuracy",
+            fig,
+        )
+
+
+# -----------------------------------------------------------------------------
+# CrossRunStabilityCallback (TEST ONLY)
+# -----------------------------------------------------------------------------
+class CrossRunStabilityCallback(BaseCallback):
+    """
+    TEST ONLY.
+
+    Stability between current dictionary D0 and checkpoint dictionary Dk using Hungarian matching
+    over cosine similarities (max average matched similarity).
+
+    Optionally also computes stability restricted to top-K concepts by test energy.
+
+    Args:
+        checkpoint_paths: list of s3://... or local paths
+        topk_by_energy: if not None, compute additional stability on top-K energetic concepts
+        n_curve_points: number of points to sample for stability vs K plots
+    """
+
+    def __init__(
+        self,
+        checkpoint_paths: List[str],
+        every_n_steps: int = None,
+        every_n_epochs: int = 1,
+        topk_by_energy: Optional[int] = 512,
+        prefix: str = "stability",
+        n_curve_points: int = 50,
+    ):
+        super().__init__(every_n_steps=every_n_steps, every_n_epochs=every_n_epochs)
+        self.checkpoint_paths = checkpoint_paths
+        self.topk_by_energy = topk_by_energy
+        self.prefix = prefix
+        self.n_curve_points = n_curve_points
+
+    def on_validation_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
+        return  # explicitly disabled
+
+    def on_test_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
+        self._compute(trainer, pl_module)
+
+    @torch.no_grad()
+    def _compute(self, trainer: Trainer, pl_module: LightningModule):
+        if not self.checkpoint_paths:
+            log.warning("[CrossRunStabilityCallback] No checkpoint paths provided.")
+            return
+
+        try:
+            from scipy.optimize import linear_sum_assignment
+        except Exception as e:
+            raise ImportError(
+                "CrossRunStabilityCallback requires scipy (scipy.optimize.linear_sum_assignment). "
+                f"Install scipy or replace Hungarian implementation. Original error: {e}"
+            )
+
+        device = pl_module.device
+
+        # Current dictionary
+        D0 = get_dictionary_from_lightningsae(pl_module).to(device)
+        D0 = F.normalize(D0, dim=-1)  # [C, d]
+        n_concepts = D0.shape[0]
+
+        # Get energy for ranking concepts
+        E = None
+        if hasattr(pl_module, "_last_energy") and pl_module._last_energy.get("combined") is not None:
+            E = pl_module._last_energy["combined"].to(device)
+
+        # Otherwise compute from test activations
+        if E is None:
+            all_test_acts = getattr(pl_module, "test_activations", None)
+            if all_test_acts:
+                for dataloader_idx, acts in all_test_acts.items():
+                    if acts.get("audio") and acts.get("text"):
+                        if isinstance(acts["audio"], list):
+                            Za = torch.cat([a.to(device) for a in acts["audio"]], dim=0)
+                        else:
+                            Za = acts["audio"].to(device)
+                        
+                        if isinstance(acts["text"], list):
+                            Zt = torch.cat([a.to(device) for a in acts["text"]], dim=0)
+                        else:
+                            Zt = acts["text"].to(device)
+                        
+                        Za = gather_tensor_if_distributed(Za, trainer)
+                        Zt = gather_tensor_if_distributed(Zt, trainer)
+                        E = 0.5 * (Za.mean(0) + Zt.mean(0))
+                        break
+
+        # Get sorted indices by energy (descending)
+        if E is not None:
+            sorted_indices = torch.argsort(E, descending=True)
+            E_sorted = E[sorted_indices]
+            total_energy = E_sorted.sum().clamp_min(1e-12)
+            cumulative_energy = torch.cumsum(E_sorted, dim=0) / total_energy
+        else:
+            sorted_indices = torch.arange(n_concepts, device=device)
+            cumulative_energy = None
+
+        # Define K values to sample for the stability curves
+        k_values = self._get_k_values(n_concepts)
+
+        def stability_hungarian(A: torch.Tensor, B: torch.Tensor) -> float:
+            sim = (A @ B.t()).detach().cpu().numpy()
+            cost = -sim
+            r, c = linear_sum_assignment(cost)
+            return float(sim[r, c].mean())
+
+        stabs_all = []
+        stabs_top = []
+        # For curve plots: stability_per_k[checkpoint_idx] = list of stabilities for each k
+        stability_per_k_all = []
+
+        for p in self.checkpoint_paths:
+            m = copy.deepcopy(pl_module).to("cpu")
+            sd = load_state_dict_any(p, map_location="cpu")
+            m.load_state_dict(sd, strict=False)
+
+            Dk = get_dictionary_from_lightningsae(m).to(device)
+            Dk = F.normalize(Dk, dim=-1)
+
+            # Full stability
+            stab = stability_hungarian(D0, Dk)
+            stabs_all.append(stab)
+            pl_module.log(f"test/{self.prefix}_vs_{_safe_name(p)}", stab, sync_dist=True)
+
+            # Top-K stability (single value)
+            if self.topk_by_energy is not None and E is not None:
+                k = min(self.topk_by_energy, n_concepts)
+                top_idx = sorted_indices[:k]
+                stab_top = stability_hungarian(D0[top_idx], Dk[top_idx])
+                stabs_top.append(stab_top)
+                pl_module.log(
+                    f"test/{self.prefix}_top{k}_vs_{_safe_name(p)}",
+                    stab_top,
+                    sync_dist=True,
+                )
+
+            # Compute stability for each K value (for curve plots)
+            stabs_for_k = []
+            for k in k_values:
+                top_idx_k = sorted_indices[:k]
+                stab_k = stability_hungarian(D0[top_idx_k], Dk[top_idx_k])
+                stabs_for_k.append(stab_k)
+            stability_per_k_all.append(stabs_for_k)
+
+        # Log mean stability
+        pl_module.log(
+            f"test/{self.prefix}_mean",
+            float(sum(stabs_all) / max(len(stabs_all), 1)),
+            prog_bar=True,
+            sync_dist=True,
+        )
+        if stabs_top:
+            k = min(self.topk_by_energy, n_concepts)
+            pl_module.log(
+                f"test/{self.prefix}_top{k}_mean",
+                float(sum(stabs_top) / max(len(stabs_top), 1)),
+                prog_bar=True,
+                sync_dist=True,
+            )
+
+        # Compute mean stability across checkpoints for each K
+        if stability_per_k_all:
+            import numpy as np
+            stab_arr = np.array(stability_per_k_all)  # [n_checkpoints, n_k_values]
+            mean_stab_per_k = stab_arr.mean(axis=0).tolist()
+
+            # Plot 1: Stability vs Number of Concepts
+            self._log_stability_vs_num_concepts_plot(
+                trainer, k_values, mean_stab_per_k
+            )
+
+            # Plot 2: Stability vs Cumulative Energy (energy-weighted)
+            if cumulative_energy is not None:
+                cum_energy_at_k = [cumulative_energy[k - 1].item() for k in k_values]
+                self._log_stability_vs_cumulative_energy_plot(
+                    trainer, cum_energy_at_k, mean_stab_per_k
+                )
+
+    def _get_k_values(self, n_concepts: int) -> List[int]:
+        """Generate K values to sample for stability curves (log-spaced for better resolution at low K)."""
+        import numpy as np
+        # Use log spacing with more points at lower K values
+        n_points = min(self.n_curve_points, n_concepts)
+        k_vals = np.unique(np.geomspace(1, n_concepts, n_points).astype(int))
+        # Ensure we have at least 1 and n_concepts
+        k_vals = sorted(set([1] + k_vals.tolist() + [n_concepts]))
+        return k_vals
+
+    def _log_stability_vs_num_concepts_plot(
+        self,
+        trainer: Trainer,
+        k_values: List[int],
+        mean_stabilities: List[float],
+    ):
+        """Plot stability as a function of number of concepts (ranked by energy)."""
+        try:
+            import plotly.graph_objects as go
+        except ImportError:
+            log.debug("plotly not available, skipping stability vs num concepts plot")
+            return
+
+        fig = go.Figure()
+        fig.add_trace(go.Scatter(
+            x=k_values,
+            y=mean_stabilities,
+            mode="lines+markers",
+            name="Mean Stability",
+            line=dict(color="green", width=2),
+            marker=dict(size=4),
+        ))
+        fig.update_layout(
+            title="Stability vs Number of Concepts (ranked by energy)",
+            xaxis_title="Number of Top Concepts (by energy)",
+            yaxis_title="Stability (Hungarian matching)",
+            yaxis_range=[0, 1.05],
+            xaxis_type="log",
+            template="plotly_white",
+        )
+
+        plot_key = f"test/{self.prefix}_vs_num_concepts_plot"
+        _log_plotly_to_wandb(trainer, plot_key, fig)
+
+    def _log_stability_vs_cumulative_energy_plot(
+        self,
+        trainer: Trainer,
+        cumulative_energy: List[float],
+        mean_stabilities: List[float],
+    ):
+        """Plot stability as a function of cumulative energy fraction (energy-weighted view)."""
+        try:
+            import plotly.graph_objects as go
+        except ImportError:
+            log.debug("plotly not available, skipping stability vs cumulative energy plot")
+            return
+
+        fig = go.Figure()
+        fig.add_trace(go.Scatter(
+            x=cumulative_energy,
+            y=mean_stabilities,
+            mode="lines+markers",
+            name="Mean Stability",
+            line=dict(color="purple", width=2),
+            marker=dict(size=4),
+        ))
+        fig.update_layout(
+            title="Stability vs Cumulative Energy Fraction",
+            xaxis_title="Cumulative Energy Fraction",
+            yaxis_title="Stability (Hungarian matching)",
+            xaxis_range=[0, 1.05],
+            yaxis_range=[0, 1.05],
+            template="plotly_white",
+        )
+
+        plot_key = f"test/{self.prefix}_vs_cumulative_energy_plot"
+        _log_plotly_to_wandb(trainer, plot_key, fig)
+
+
+# -----------------------------------------------------------------------------
+# SimLoggerCallback
+# -----------------------------------------------------------------------------
+class SimLoggerCallback(BaseCallback):
+    """
+    Logs distributions of pairwise cosine similarities between embeddings:
+      - audio-audio (blue)
+      - text-text (red)
+      - text-audio off-diagonal (purple)
+      - text-audio diagonal / matched pairs (dark purple)
+
+    Reads embeddings from pl_module.val_embeddings / pl_module.test_embeddings
+    (populated by LightningSAE._eval_step).
+    """
+
+    def __init__(
+        self,
+        enable_on_validation: bool = True,
+        enable_on_test: bool = True,
+        every_n_steps: int = None,
+        every_n_epochs: int = 1,
+        max_samples: int = 2048,
+        n_bins: int = 50,
+        prefix: str = "similarity",
+    ):
+        super().__init__(every_n_steps=every_n_steps, every_n_epochs=every_n_epochs)
+        self.enable_on_validation = enable_on_validation
+        self.enable_on_test = enable_on_test
+        self.max_samples = max_samples
+        self.n_bins = n_bins
+        self.prefix = prefix
+
+    def on_validation_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
+        if not self.enable_on_validation:
+            return
+        if not (self._check_step(trainer, pl_module) or self._check_epoch(trainer, pl_module)):
+            return
+        self._compute_and_log(trainer, pl_module, mode="val")
+
+    def on_test_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
+        if not self.enable_on_test:
+            return
+        if not (self._check_step(trainer, pl_module) or self._check_epoch(trainer, pl_module)):
+            return
+        self._compute_and_log(trainer, pl_module, mode="test")
+
+    @staticmethod
+    @torch.no_grad()
+    def _compute(
+        audio_emb: Optional[torch.Tensor],
+        text_emb: Optional[torch.Tensor],
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Pure computation — no trainer / pl_module / logging.
+
+        Args:
+            audio_emb: [N_a, D] or None
+            text_emb:  [N_t, D] or None
+
+        Returns:
+            dict mapping label -> 1D tensor of cosine similarities.
+            Labels: 'audio-audio', 'text-text', 'text-audio (matched)',
+                    'text-audio (off-diag)'.
+        """
+        def cosine_sim_matrix(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+            return F.normalize(a, dim=-1) @ F.normalize(b, dim=-1).T
+
+        sims = {}
+
+        if audio_emb is not None and audio_emb.size(0) > 1:
+            sim_aa = cosine_sim_matrix(audio_emb, audio_emb)
+            triu_idx = torch.triu_indices(sim_aa.size(0), sim_aa.size(1), offset=1)
+            sims["audio-audio"] = sim_aa[triu_idx[0], triu_idx[1]]
+
+        if text_emb is not None and text_emb.size(0) > 1:
+            sim_tt = cosine_sim_matrix(text_emb, text_emb)
+            triu_idx = torch.triu_indices(sim_tt.size(0), sim_tt.size(1), offset=1)
+            sims["text-text"] = sim_tt[triu_idx[0], triu_idx[1]]
+
+        if audio_emb is not None and text_emb is not None:
+            sim_ta = cosine_sim_matrix(text_emb, audio_emb)
+            n_diag = min(sim_ta.shape[0], sim_ta.shape[1])
+            sims["text-audio (matched)"] = torch.diag(sim_ta[:n_diag, :n_diag])
+            mask = torch.ones_like(sim_ta, dtype=torch.bool)
+            for i in range(n_diag):
+                mask[i, i] = False
+            sims["text-audio (off-diag)"] = sim_ta[mask]
+
+        return sims
+
+    @torch.no_grad()
+    def _compute_and_log(self, trainer: Trainer, pl_module: LightningModule, mode: str):
+        """Gather data from module, call _compute, then log."""
+        all_embs = getattr(pl_module, f"{mode}_embeddings", {})
+        device = pl_module.device
+
+        for dataloader_idx, embs in all_embs.items():
+            audio_emb = _resolve_modality_tensors(embs, "audio", device, trainer)
+            text_emb = _resolve_modality_tensors(embs, "text", device, trainer)
+
+            if audio_emb is None and text_emb is None:
+                log.warning(f"[SimLoggerCallback] No embeddings for mode={mode}, dataloader_idx={dataloader_idx}.")
+                continue
+
+            # Move to CPU and subsample
+            if audio_emb is not None:
+                audio_emb = audio_emb.cpu()
+                if audio_emb.size(0) > self.max_samples:
+                    audio_emb = audio_emb[torch.randperm(audio_emb.size(0))[: self.max_samples]]
+            if text_emb is not None:
+                text_emb = text_emb.cpu()
+                if text_emb.size(0) > self.max_samples:
+                    text_emb = text_emb[torch.randperm(text_emb.size(0))[: self.max_samples]]
+
+            sims = self._compute(audio_emb, text_emb)
+            if not sims:
+                continue
+
+            dataset_name = _get_dataset_name(trainer, dataloader_idx, mode)
+            self._log_similarity_figure(trainer, sims, mode, dataset_name)
+
+    def _log_similarity_figure(
+        self,
+        trainer: Trainer,
+        sims: Dict[str, torch.Tensor],
+        mode: str,
+        dataset_name: str,
+    ):
+        """Build and log plotly figure from pre-computed similarity tensors."""
+        try:
+            import plotly.graph_objects as go
+        except ImportError:
+            log.debug("plotly not available, skipping similarity distribution plot")
+            return
+
+        colors = {
+            "audio-audio": "rgba(0, 0, 255, 0.6)",
+            "text-text": "rgba(255, 0, 0, 0.6)",
+            "text-audio (off-diag)": "rgba(128, 0, 128, 0.5)",
+            "text-audio (matched)": "rgba(75, 0, 130, 0.8)",
+        }
+
+        fig = go.Figure()
+        for label, vals in sims.items():
+            fig.add_trace(go.Histogram(
+                x=vals.cpu().numpy().tolist(),
+                name=label,
+                opacity=0.7,
+                marker_color=colors.get(label, "gray"),
+                nbinsx=self.n_bins,
+                histnorm="probability density",
+            ))
+
+        fig.update_layout(
+            title=f"Embedding Cosine Similarity Distributions ({mode})",
+            xaxis_title="Cosine Similarity",
+            yaxis_title="Density",
+            barmode="overlay",
+            template="plotly_white",
+            legend=dict(yanchor="top", y=0.99, xanchor="left", x=0.01),
+        )
+
+        _log_plotly_to_wandb(trainer, f"Similarity/{dataset_name}/distributions", fig)

steerable_retrieval/callbacks/save.py

@@ -0,0 +1,129 @@
+"""
+SaveActivationsCallback — saves validation embeddings, preactivations,
+activations, and reconstructions to per-tensor .pt files at the end of each
+validation epoch.
+
+Reads embeddings and activations stored on the LightningModule by _eval_step,
+then re-runs the SAE forward pass to capture preactivations and reconstructions
+(which are not stored on the module to save memory during normal training).
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Dict, Optional
+
+import torch
+from lightning.pytorch import Trainer
+from lightning.pytorch.core import LightningModule
+
+from steerable_retrieval.callbacks.utils import BaseCallback
+from steerable_retrieval.callbacks.energy import _get_dataset_name, gather_tensor_if_distributed, _cat_or_none
+
+log = logging.getLogger(__name__)
+
+
+class SaveActivationsCallback(BaseCallback):
+    """
+    At the end of every validation epoch, saves separate .pt files containing
+    embeddings, preactivations, activations, and reconstructions for every
+    validation dataloader / modality.
+
+    File structure::
+
+        <save_dir>/activations/<dataset_name>/<modality>/<tensor_name>/tensors.pt
+
+    Example::
+
+        activations/dataset_name/audio/embeddings/tensors.pt
+        activations/dataset_name/audio/preactivations/tensors.pt
+        activations/dataset_name/audio/activations/tensors.pt
+        activations/dataset_name/audio/reconstructions/tensors.pt
+    """
+
+    ROOT_DIRNAME = "activations"
+
+    def __init__(
+        self,
+        save_dir: str,
+        every_n_steps: int = None,
+        every_n_epochs: int = 1,
+    ):
+        super().__init__(every_n_steps=every_n_steps, every_n_epochs=every_n_epochs)
+        self.save_dir = save_dir
+
+    def on_validation_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
+        if not (self._check_step(trainer, pl_module) or self._check_epoch(trainer, pl_module)):
+            return
+        self._save(trainer, pl_module)
+
+    @torch.no_grad()
+    def _save(self, trainer: Trainer, pl_module: LightningModule):
+        """Build tensors from module-stored embeddings/activations and save to disk."""
+        device = pl_module.device
+        all_acts = getattr(pl_module, "val_activations", {})
+        all_embs = getattr(pl_module, "val_embeddings", {})
+
+        for dataloader_idx in all_acts:
+            dataset_name = _get_dataset_name(trainer, dataloader_idx, "val")
+
+            acts_dl = all_acts.get(dataloader_idx, {})
+            embs_dl = all_embs.get(dataloader_idx, {})
+
+            for modality in ("audio", "text"):
+                act_chunks = acts_dl.get(modality, [])
+                emb_chunks = embs_dl.get(modality, [])
+                if not act_chunks:
+                    continue
+
+                # Concatenate and gather
+                activations = _cat_or_none(act_chunks, device=device)
+                embeddings = _cat_or_none(emb_chunks, device=device) if emb_chunks else None
+
+                activations = gather_tensor_if_distributed(activations, trainer)
+                if embeddings is not None:
+                    embeddings = gather_tensor_if_distributed(embeddings, trainer)
+
+                # Re-run SAE forward on embeddings to get preactivations & reconstructions
+                preactivations = None
+                reconstructions = None
+                if embeddings is not None:
+                    # Process in chunks to avoid OOM
+                    chunk_size = 512
+                    pre_chunks, rec_chunks = [], []
+                    for i in range(0, embeddings.size(0), chunk_size):
+                        emb_chunk = embeddings[i : i + chunk_size].to(device)
+                        _, z, xhat, pre = pl_module(emb_chunk)
+                        pre_chunks.append(pre.cpu() if pre is not None else torch.zeros_like(z).cpu())
+                        rec_chunks.append(xhat.cpu())
+                    preactivations = torch.cat(pre_chunks, dim=0)
+                    reconstructions = torch.cat(rec_chunks, dim=0)
+
+                tensors = {
+                    "embeddings": embeddings.cpu() if embeddings is not None else None,
+                    "preactivations": preactivations,
+                    "activations": activations.cpu(),
+                    "reconstructions": reconstructions,
+                }
+                if trainer.is_global_zero:
+                    self._save_tensors(dataset_name=dataset_name, modality=modality, tensors=tensors)
+
+        # Only rank-0 writes to disk
+        if trainer.is_global_zero:
+            log.info(f"[SaveActivationsCallback] Saved validation activations to {self.save_dir}")
+
+    def _save_tensors(
+        self,
+        dataset_name: str,
+        modality: str,
+        tensors: Dict[str, Optional[torch.Tensor]],
+    ) -> None:
+        base_dir = os.path.join(self.save_dir, self.ROOT_DIRNAME, dataset_name, modality)
+        for tensor_name, tensor_value in tensors.items():
+            if tensor_value is None:
+                continue
+            tensor_dir = os.path.join(base_dir, tensor_name)
+            os.makedirs(tensor_dir, exist_ok=True)
+            tensor_path = os.path.join(tensor_dir, "tensors.pt")
+            torch.save(tensor_value, tensor_path)

steerable_retrieval/callbacks/utils.py

@@ -0,0 +1,36 @@
+from lightning.pytorch.callbacks import Callback
+from lightning.pytorch import Trainer
+from lightning.pytorch.core import LightningModule
+
+import torch
+
+
+class BaseCallback(Callback):
+
+    def __init__(self, every_n_steps = 1, every_n_epochs = 1, **kwargs):
+        super().__init__(**kwargs)
+        self.every_n_steps = every_n_steps
+        self.every_n_epochs = every_n_epochs
+        
+    def _check_step(self, trainer: Trainer, pl_module: LightningModule) -> bool:
+        if self.every_n_steps is not None:
+            return trainer.global_step % self.every_n_steps == 0
+        else:
+            return False
+        
+    def _check_epoch(self, trainer: Trainer, pl_module: LightningModule) -> bool:
+        if self.every_n_epochs is not None:
+            return trainer.current_epoch % self.every_n_epochs == 0
+        else:
+            return False
+
+    def _should_run(self, trainer: Trainer, pl_module: LightningModule) -> bool:
+        return self._check_step(trainer, pl_module) or self._check_epoch(trainer, pl_module)
+
+    def _should_run_on_validation(self, trainer: Trainer, pl_module: LightningModule) -> bool:
+        return self._check_step(trainer, pl_module) or self._check_epoch(trainer, pl_module)
+
+    def _should_run_on_test(self, trainer: Trainer, pl_module: LightningModule) -> bool:
+        return self._check_step(trainer, pl_module) or self._check_epoch(trainer, pl_module)
+
+

steerable_retrieval/dataloading/__init__.py

@@ -0,0 +1 @@
+

steerable_retrieval/dataloading/dataloaders.py

@@ -0,0 +1,319 @@
+import logging
+import os
+import random
+from hashlib import sha256
+
+import pandas as pd
+from lightning import LightningDataModule
+from sklearn.model_selection import train_test_split
+from torch.utils.data import DataLoader
+
+from .datasets import TextAudioDataset
+from steerable_retrieval.utils.instantiators import instantiate
+
+
+def instantiate_datasets(datasets_cfg):
+    datasets = []
+    dataset_names = []
+    for dataset_name in datasets_cfg:
+        dataset = instantiate(datasets_cfg[dataset_name])
+        datasets.append(dataset)
+        dataset_names.append(dataset_name)
+    return datasets, dataset_names
+
+
+def get_song_describer_annotations(data_path = None, csv_path = None, val_split = 0.1):
+    
+    df = pd.read_csv(csv_path)
+    
+    
+    df = df[['path','caption','is_valid_subset','caption_id']].rename(columns = {'path':'file_path'})
+    df['file_path'] = os.path.join(data_path) + '/' + df['file_path']
+    #replace .mp3 with .2min.mp3
+    df['file_path'] = df['file_path'].apply(lambda x: x.replace('.mp3','.2min.mp3'))
+    
+    records = df.to_dict(orient = 'records')
+    
+    for record in records:
+        record['caption'] = {sha256(record['caption'].encode('utf-8')).hexdigest(): record['caption']}
+    if val_split == 0.0:
+        print('No validation split')
+        for record in records:
+            record['split'] = 'train'
+        return records
+    
+    train_indices, val_indices = train_test_split(range(len(records)), test_size = val_split, random_state = 42)
+    
+    for idx in train_indices:
+        records[idx]['split'] = 'train'
+        
+    for idx in val_indices:
+        records[idx]['split'] = 'val'
+    
+    return records
+
+
+
+def get_musiccaps_annotations(data_path = None, csv_path = None, val_split = 0.1, test_split = 0.1):
+        
+        df = pd.read_csv(csv_path)
+        df['file_path'] = data_path + '/' + df['ytid'] + '.wav'
+        
+        records = df.to_dict(orient = 'records')
+        
+        for record in records:
+            record['caption'] = {sha256(record['caption'].encode('utf-8')).hexdigest(): record['caption']}
+            
+        if val_split == 0.0:
+            print('No validation split')
+            for record in records:
+                record['split'] = 'train'
+            return records
+        
+        # split into train, val, test
+        train_indices, test_indices = train_test_split(range(len(records)), test_size = test_split + val_split, random_state = 42)
+        val_indices, test_indices = train_test_split(test_indices, test_size = test_split/(test_split + val_split), random_state = 42)
+        
+        
+        
+        for idx in train_indices:
+            records[idx]['split'] = 'train'
+            
+        for idx in val_indices:
+            records[idx]['split'] = 'val'
+            
+        for idx in test_indices:
+            records[idx]['split'] = 'test'
+            
+        return records
+    
+def get_musiccaps_truncated_annotations(data_path = None, csv_path = None, val_split = 0.1, test_split = 0.1):
+    df = pd.read_csv(csv_path)
+    df['file_path'] = data_path + '/' + df['ytid'] + '.wav'
+    
+    import random
+    
+    records = df.to_dict(orient = 'records')
+    
+    print('Truncating captions')
+    
+    for record in records:
+        # select a random sentence from the caption
+        sentences = record['caption'].split('.')
+        random_sentence = random.choice(sentences)
+        record['caption'] = {sha256(random_sentence.encode('utf-8')).hexdigest(): random_sentence}
+        
+    if val_split == 0.0:
+        print('No validation split')
+        for record in records:
+            record['split'] = 'train'
+        return records
+    
+    # split into train, val, test
+    train_indices, test_indices = train_test_split(range(len(records)), test_size = test_split + val_split, random_state = 42)
+    val_indices, test_indices = train_test_split(test_indices, test_size = test_split/(test_split + val_split), random_state = 42)
+    
+    for idx in train_indices:
+        records[idx]['split'] = 'train'
+        
+    for idx in val_indices:
+        records[idx]['split'] = 'val'
+        
+    for idx in test_indices:
+        records[idx]['split'] = 'test'
+        
+    return records
+
+
+def get_maxcaps_annotations(data_path = None, csv_path = None):
+    """Read JSONL file line-by-line to avoid memory issues with large files."""
+
+    df = pd.read_csv(csv_path)
+    df['file_path'] = data_path + '/' + df['file_path']
+    records = df.to_dict(orient = 'records')
+    for record in records:
+        record['caption'] = {sha256(record['caption'].encode('utf-8')).hexdigest(): record['caption']}
+    return records
+
+def get_yt8m_annotations(data_path = None, csv_path = None):
+    df = pd.read_csv(csv_path)
+    df['file_path'] = data_path + '/' + df['file_path']
+    records = df.to_dict(orient = 'records')
+    for record in records:
+        record['caption'] = {sha256(record['caption'].encode('utf-8')).hexdigest(): record['caption']}
+    return records
+
+def get_music4all_annotations(manifest_csv=None, audio_status_ok=True):
+    """Annotations for music4all from the MuQ-MuLan embedding manifest.
+
+    ``file_path`` is set directly to the pre-extracted ``.npy`` audio-embedding path so
+    the preextracted-feature loader reads it as-is. Splits (train/val/test) and captions
+    are taken from the manifest. Only rows with a successfully embedded audio are kept.
+    """
+    df = pd.read_csv(manifest_csv)
+    if audio_status_ok and 'audio_embedding_status' in df.columns:
+        df = df[df['audio_embedding_status'] == 'ok']
+    df = df[df['audio_embedding_path'].notna()]
+
+    records = []
+    for row in df.itertuples(index=False):
+        caption = str(getattr(row, 'caption', '') or '')
+        split = getattr(row, 'split', 'train')
+        records.append({
+            'file_path': row.audio_embedding_path,
+            'caption': {sha256(caption.encode('utf-8')).hexdigest(): caption},
+            'split': split if split in ('train', 'val', 'test') else 'train',
+            'track_id': getattr(row, 'track_id', None),
+        })
+    logging.info(f"Loaded {len(records)} music4all annotations from {manifest_csv}")
+    return records
+
+
+def get_folder_annotations(data_path = None):
+    
+    # recursively get all files in the data_path directory that are audio files, and their paths
+    audio_files = []
+    
+    for root, dirs, files in os.walk(data_path):
+        audio_files += [os.path.join(root, file) for file in files if file.endswith('.wav') or file.endswith('.mp3')]
+        
+    records = [{'file_path': file, 'caption': '', 'split': 'train'} for file in audio_files]
+    
+    logging.info(f"Found {len(records)} audio files in {data_path}")
+
+
+    return records
+    
+    
+
+class TextAudioDataModule(LightningDataModule):
+    
+    def __init__(self,
+    datasets,
+    return_audio = True,
+    return_text = True,
+    concept = None, 
+    target_n_samples = 96000, 
+    target_sr = 48000, 
+    batch_size = 32, num_workers = 0, preextracted_features = False, truncate_preextracted = 50, root_dir = None, new_dir = None,
+    **kwargs):
+
+
+        super().__init__()
+        self.annotations = []
+        dataset_names, datasets = list(datasets.keys()), list(datasets.values())
+
+
+
+
+        self.datasets = datasets
+        self.dataset_names = dataset_names
+        for dataset in self.datasets:
+            dataset.split = dataset.split
+
+        self.return_audio = return_audio
+        self.return_text = return_text
+        self.concept = concept
+        self.target_n_samples = target_n_samples
+        self.target_sr = target_sr
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.preextracted_features = preextracted_features
+        self.truncate_preextracted = truncate_preextracted
+        
+        
+        self.truncate_preextracted = [self.truncate_preextracted for _ in range(len(self.datasets))]
+        
+        self.root_dirs = [dataset.root_dir for dataset in self.datasets]
+        self.new_dirs = [dataset.new_dir for dataset in self.datasets]
+        
+        self.train_annotations = []
+        self.val_annotations = []
+        self.test_annotations = []
+        self.val_dataset_names = []
+        self.test_dataset_names = []
+        self.val_dataset_indices = []  # Track original dataset indices for val
+        self.test_dataset_indices = []  # Track original dataset indices for test
+        
+        
+        for i, dataset_ in enumerate(self.datasets):
+            self.train_annotations.extend([annot for annot in dataset_.annotations if annot['split'] == 'train'])
+            val_annots = [annot for annot in dataset_.annotations if annot['split'] == 'val']
+            test_annots = [annot for annot in dataset_.annotations if annot['split'] == 'test']
+            if val_annots:
+                self.val_annotations.append(val_annots)
+                self.val_dataset_names.append(self.dataset_names[i])
+                self.val_dataset_indices.append(i)
+            if test_annots:
+                self.test_annotations.append(test_annots)
+                self.test_dataset_names.append(self.dataset_names[i])
+                self.test_dataset_indices.append(i)
+        
+
+        # filter for empty validation and test annotations (already filtered above, but keeping for safety)
+        self.val_annotations = [annot for annot in self.val_annotations if annot]
+        self.test_annotations = [annot for annot in self.test_annotations if annot]
+        
+        # Create dataloader_names mapping: maps dataloader_idx -> dataset_name (for val and test)
+        # This will be used by callbacks to properly name metrics
+        self.dataloader_names = {}
+        for idx, name in enumerate(self.val_dataset_names):
+            self.dataloader_names[idx] = name
+        self.test_dataloader_names = {}
+        for idx, name in enumerate(self.test_dataset_names):
+            self.test_dataloader_names[idx] = name
+
+        print(f"Number of training samples: {len(self.train_annotations)}")
+        print(f"Number of validation samples: {sum([len(annot) for annot in self.val_annotations])} over {len(self.val_annotations)} datasets, {[len(annot) for annot in self.val_annotations]}")
+        print(f"Number of test samples: {sum([len(annot) for annot in self.test_annotations])} over {len(self.test_annotations)} datasets, {[len(annot) for annot in self.test_annotations]}")
+        print(f"Datasets: {self.dataset_names}")
+
+    @property
+    def names(self):
+        """Mapping mode -> (dataloader_idx -> dataset name). Used by callbacks for log keys."""
+        return {
+            "val": getattr(self, "dataloader_names", None) or {},
+            "test": getattr(self, "test_dataloader_names", None) or {},
+        }
+
+    def setup(self, stage: str) -> None:
+        
+        if stage != 'eval':
+            self.train_dataset = TextAudioDataset(annotations=self.train_annotations, target_n_samples=self.target_n_samples, target_sr=self.target_sr, return_audio=self.return_audio, return_text=self.return_text, concept=self.concept, preextracted_features=self.preextracted_features, truncate_preextracted=self.truncate_preextracted[0], root_dir=self.root_dirs[0], new_dir=self.new_dirs[0]) 
+            
+        self.val_datasets = [TextAudioDataset(annotations=self.val_annotations[i], target_n_samples=self.target_n_samples, target_sr=self.target_sr, return_audio=self.return_audio, return_text=self.return_text, concept=self.concept, preextracted_features=self.preextracted_features, truncate_preextracted=self.truncate_preextracted[self.val_dataset_indices[i]], root_dir=self.root_dirs[self.val_dataset_indices[i]], new_dir=self.new_dirs[self.val_dataset_indices[i]]) for i in range(len(self.val_annotations))]
+        self.test_datasets = [TextAudioDataset(annotations=self.test_annotations[i], target_n_samples=self.target_n_samples, target_sr=self.target_sr, return_audio=self.return_audio, return_text=self.return_text, concept=self.concept, preextracted_features=self.preextracted_features, truncate_preextracted=self.truncate_preextracted[self.test_dataset_indices[i]], root_dir=self.root_dirs[self.test_dataset_indices[i]], new_dir=self.new_dirs[self.test_dataset_indices[i]]) for i in range(len(self.test_annotations))]
+        
+        
+    def train_dataloader(self):
+        return DataLoader(
+            self.train_dataset, 
+            batch_size=self.batch_size, 
+            num_workers=self.num_workers, 
+            shuffle=True,
+            drop_last=True  # Critical for distributed training to avoid hangs
+        )
+    
+    def val_dataloader(self):
+        return [
+            DataLoader(
+                val_dataset, 
+                batch_size=self.batch_size, 
+                num_workers=self.num_workers, 
+                shuffle=False,
+                drop_last=False  # Don't drop for validation
+            ) 
+            for val_dataset in self.val_datasets
+        ]
+    
+    def test_dataloader(self):
+        return [
+            DataLoader(
+                test_dataset, 
+                batch_size=self.batch_size, 
+                num_workers=self.num_workers, 
+                shuffle=False,
+                drop_last=False  # Don't drop for test
+            ) 
+            for test_dataset in self.test_datasets
+        ]

steerable_retrieval/dataloading/datasets.py

@@ -0,0 +1,437 @@
+from torch.utils.data import Dataset
+from .loading_utils import *
+import torch
+import os
+import random
+from tqdm import tqdm
+import pandas as pd
+from hydra.utils import instantiate
+
+import logging
+
+class BasicProcessor:
+    def __init__(self, probability = 1.0, split = None):
+        self.probability = probability
+        self.split = split
+    
+    def process(self, annot):
+        """
+        Process the annotation if the probability is less than the probability and the split is in the split list
+        """
+        if random.random() < self.probability:
+            if self.split is None:
+                return self.process(annot)
+            if self.split is not None and annot['split'] in self.split:
+                return self.process(annot)
+            else:
+                return annot
+        else:
+            return annot
+        
+    def __call__(self, annot):
+        return self.process(annot)
+    
+class RandomNSentencesProcessor(BasicProcessor):
+
+    def process(self, annot):
+        
+        prompt = annot['prompt']
+        sentences = prompt.split('.')
+        n_sentences = len(sentences)
+        keep_n_sentences = random.randint(1, n_sentences)
+        random_sentences = random.sample(random_sentences, keep_n_sentences)
+        prompt = '. '.join(random_sentences)
+        return prompt
+    
+class ShuffleSentencesProcessor(BasicProcessor):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+    
+    def process(self, annot):
+        prompt = annot['prompt']
+        sentences = prompt.split('.')
+        random.shuffle(sentences)
+        prompt = '. '.join(sentences)
+        return prompt
+
+class ShuffleTagsProcessor(BasicProcessor):
+    def __init__(self, replace_caption_p = 0.5, **kwargs):
+        super().__init__(**kwargs)
+        self.replace_caption_p = replace_caption_p
+    
+    def process(self, annot):
+        ## if tags is in the annotation, it will be a comma-separated string
+        ## drop anywhere from all but 1 to 1 tag, and shuffle
+        if 'tags' in annot:
+            tags = annot['tags'].split(',')
+        else:
+            return annot
+        len_tags = len(tags)
+        drop_tags = random.randint(1, len_tags-1)
+        tags_to_drop = random.sample(tags, drop_tags)
+        tags = [tag for tag in tags if tag not in tags_to_drop]
+        random.shuffle(tags)
+        annot['tags'] = '. '.join(tags)
+
+        # if replace caption, replace else add it to the caption
+        if random.random() < self.replace_caption_p:
+            annot['prompt'] = annot['tags']
+        else:
+            annot['prompt'] = annot['prompt'] + '.' + annot['tags']
+        return annot
+
+class TextAudioDataset(Dataset):
+    def __init__(self,
+                 annotations = None,
+                 get_annotations_function = None,
+                 task_kwargs = None,
+                 target_n_samples = 96000,
+                 target_sr = 48000,
+                 return_audio = True,
+                 return_text = True,
+                 concept = None,
+                 return_full_audio = False,
+                 preextracted_features = False,
+                 truncate_preextracted = 50,
+                 split = None,
+                 filter_split = None,
+                 root_dir = None,
+                 new_dir = None,
+                 limit_n = None,
+                 processors = [],
+                 **kwargs
+                 ):
+        
+        
+        # Get annotations either directly or from function
+        if annotations is not None:
+            self.annotations = annotations
+        elif get_annotations_function is not None:
+            # Support both string and callable
+            if isinstance(get_annotations_function, str):
+                task_kwargs = task_kwargs or {}
+                import importlib
+
+                # parse the fully qualified function name
+                module_name, func_name = get_annotations_function.rsplit('.', 1)
+                module = importlib.import_module(module_name)
+                get_annotations_func = getattr(module, func_name)
+                self.annotations = get_annotations_func(**task_kwargs)
+            else:
+                task_kwargs = task_kwargs or {}
+                self.annotations = get_annotations_function(**task_kwargs)
+        else:
+            raise ValueError("Must provide either annotations or get_annotations_function")
+        
+
+
+
+        self.target_n_samples = target_n_samples
+        self.target_sr = target_sr
+        self.return_audio = return_audio
+        self.return_text = return_text
+        self.concept = concept
+        self.return_full_audio = return_full_audio
+        self.preextracted_features = preextracted_features
+        self.truncate_preextracted = truncate_preextracted
+        self.split = split
+        self.root_dir = root_dir
+        self.new_dir = new_dir
+        self.limit_n = limit_n
+        # Update split if needed
+        if split is not None and split != 'keep':
+            for annot in self.annotations:
+                annot['split'] = split
+        elif split == 'keep':
+            # Keep original splits from annotations, or set to 'train' if not present
+            for annot in self.annotations:
+                if 'split' not in annot or annot['split'] not in ['train', 'val', 'test']:
+                    annot['split'] = 'train'
+        elif split is None and len(self.annotations) > 0 and 'split' not in self.annotations[0].keys():
+            for annot in self.annotations:
+                annot['split'] = 'train'
+
+        if filter_split is not None:
+            self.annotations = [annot for annot in self.annotations if annot['split'] in filter_split]
+                
+        annot_df = pd.DataFrame(self.annotations)
+        
+        try:
+            annot_df['file_index'] = pd.factorize(annot_df['file_path'])[0]
+        except Exception as e:
+            print(e)
+        
+        annot_df['file_path'] = annot_df['file_path'].apply(lambda x: x.replace(root_dir, new_dir) if root_dir is not None and new_dir is not None else x)
+        
+        self.annotations = annot_df.to_dict('records')
+
+        # Filter out annotations whose feature/audio file is missing (avoids IndexError when retrying in __getitem__)
+        if self.return_audio and self.preextracted_features:
+            n_before = len(self.annotations)
+            self.annotations = [
+                a for a in self.annotations
+                if os.path.exists(a['file_path'].replace('.mp3', '.npy').replace('.wav', '.npy'))
+            ]
+            if len(self.annotations) < n_before:
+                logging.warning(
+                    f"Filtered out {n_before - len(self.annotations)} annotations with missing feature files. "
+                    f"Dataset has {len(self.annotations)} samples."
+                )
+        elif self.return_audio and not self.preextracted_features:
+            n_before = len(self.annotations)
+            def _audio_exists(a):
+                p = a['file_path']
+                return os.path.exists(p.replace('.npy', '.mp3')) or os.path.exists(p.replace('.npy', '.wav'))
+            self.annotations = [a for a in self.annotations if _audio_exists(a)]
+            if len(self.annotations) < n_before:
+                logging.warning(
+                    f"Filtered out {n_before - len(self.annotations)} annotations with missing audio files. "
+                    f"Dataset has {len(self.annotations)} samples."
+                )
+
+        if self.limit_n is not None and self.limit_n < len(self.annotations):
+            self.annotations = self.annotations[:self.limit_n]
+            print(f"Limiting dataset to {self.limit_n} samples")
+        else:
+            print(f"Dataset has {len(self.annotations)} samples")
+
+        
+        assert return_audio or return_text, "At least one of return_audio or return_text must be True (duh)"
+
+        self.processors = [instantiate(processor) for processor in processors]
+
+        
+
+
+    def purge(self):
+        if self.return_audio and not self.preextracted_features:
+            raise NotImplementedError("Purging your audio dataset is probably a bad idea")
+        else:
+            file_paths = [annot['file_path'] for annot in self.annotations]
+            for file_path in file_paths:
+                os.remove(file_path)
+            print(f"Removed {len(file_paths)} files")
+        
+    def __len__(self):
+        return len(self.annotations)
+
+    def __getitem__(self, idx, return_full_audio = False, hop = None, verbose = False):
+        
+        
+        return_full_audio = self.return_full_audio if return_full_audio is None else return_full_audio
+        
+        annot = self.annotations[idx]
+        
+        
+        if self.return_audio:
+            if not self.preextracted_features:
+                # file_path = annot['file_path'].replace('.npy','.mp3').replace('.wav','.mp3')
+                # check if the mp3 file exists
+                if os.path.exists(annot['file_path'].replace('.npy','.mp3')):
+                    file_path = annot['file_path'].replace('.npy','.mp3')
+                elif os.path.exists(annot['file_path'].replace('.npy','.wav')):
+                    file_path = annot['file_path'].replace('.npy','.wav')
+                else:
+                    return self.__getitem__(idx+1)
+                
+                annot['file_path'] = file_path
+                try:
+                    audio = load_full_and_split(
+                        annot['file_path'],
+                        self.target_sr,
+                        self.target_n_samples,
+                        hop=hop,
+                        verbose=verbose
+                        ) if return_full_audio else load_audio_chunk(
+                        annot['file_path'],
+                        target_sr=self.target_sr,
+                        target_n_samples=self.target_n_samples,
+                        verbose=verbose
+                        )
+                    audio = audio.mean(1)
+                except Exception as e:
+                    return self.__getitem__(idx+1)
+            else:
+                file_path = annot['file_path'].replace('.mp3','.npy').replace('.wav','.npy')
+                try:
+
+                    audio = np.load(file_path,mmap_mode='r')
+                    # Preextracted features may be stored per-clip as a single vector [D]
+                    # or as multiple frames [T, D]; sample a random frame when framewise.
+                    if audio.ndim > 1:
+                        rand_index = random.randint(0, audio.shape[0]-1)
+                        audio = audio[rand_index]
+                    audio = torch.tensor(np.asarray(audio))
+                except Exception as e:
+                    return self.__getitem__(idx+1)
+                
+                
+        
+        if self.return_text:
+            possible_captions = annot['caption']
+            # ramdomly choose a caption hash
+            
+            random_hash = random.choice(list(possible_captions.keys()))
+        
+            caption = possible_captions[random_hash]
+                    
+        return_dict = {}
+        
+        if self.return_audio:
+            # Clone so storage is resizable; avoids DataLoader collate error with mmap/numpy-derived tensors
+            return_dict['audio'] = audio.clone() if isinstance(audio, torch.Tensor) else audio
+            return_dict['file_path'] = annot['file_path']
+            
+        if self.return_text:
+            return_dict['prompt'] = caption
+                
+        return_dict['file_idx'] = annot['file_index']
+
+        for processor in self.processors:
+            return_dict = processor(return_dict)
+            
+        return return_dict
+    
+    
+    def extract_features(self, model, extract_method = 'extract_features', extract_kwargs = {}, out_key = 'embedding',hop = None, return_full_audio = True, verbose = False):
+        
+        device = next(model.parameters()).device
+        print(f"Extracting features with {extract_method} method on {device} device") if verbose else None
+        
+        
+        for param in model.parameters():
+            param.requires_grad = False
+        try:
+            model.eval()
+        except:
+            pass
+        
+        for i in range(len(self)):
+            try:
+                item = self.__getitem__(i, return_full_audio = return_full_audio, hop = hop, verbose = verbose)
+                file_path = self.annotations[i]['file_path'].replace('.mp3','.npy').replace('.wav','.npy')
+                
+                audio = item['audio'].squeeze(1).to(device)
+
+
+                
+                if audio.shape[0] > 200 :
+                    chunks = torch.split(audio, 200, dim=0)
+                    chunks = list(chunks)
+                    audio_features = []
+                    for chunk in chunks:
+                        feat = getattr(model, extract_method)(chunk, **extract_kwargs)
+                        if out_key is not None:
+                            feat = feat[out_key]
+                        audio_features.append(feat)
+                    audio_features = torch.cat(audio_features, dim=0)
+                else:
+                    audio_features = getattr(model, extract_method)(audio.to(device), **extract_kwargs)
+                    if out_key is not None:
+                        audio_features = audio_features[out_key]
+                
+                    
+                print(f"Extracted features for {file_path}, shape: {audio_features.shape}") if verbose else None
+                
+                yield audio_features, file_path
+            except Exception as e:
+                print(f"Error extracting features for {file_path}: {e}") if verbose else None
+                continue
+            
+    def extract_and_save_features(self, model, save_dir = None, extract_method = 'extract_features', extract_kwargs = {}, out_key = 'embedding', hop = None, return_full_audio = True, limit_n = None, save = False, verbose = True, root_path = None, done_ids = None):
+        
+        
+        print(self.__len__())
+        
+        audio_features_all = []
+        counter = 0
+        skipped_count = 0
+        
+        save_dir = '' if save_dir is None else save_dir
+        done_ids = done_ids or set()
+        
+        if 's3://' in save_dir:
+            import boto3
+            import io
+            client = boto3.client('s3')
+        else:
+            client = None
+            import io
+
+        # filter self.annotations to only include files that are not in done_ids
+        new_annotations = []
+        for annot in self.annotations:
+            fp = annot['file_path']
+            fp = fp.replace(root_path+'/','')
+            # remove extension
+            fp = fp.replace('.mp3','').replace('.wav','').replace('.npy','')
+
+            if fp not in done_ids:
+                new_annotations.append(annot)
+        
+        self.annotations = new_annotations
+        
+        for audio_features, file_path in (pbar:= tqdm(self.extract_features(model, extract_method = extract_method, extract_kwargs = extract_kwargs, out_key = out_key, hop = hop, return_full_audio = return_full_audio, verbose = verbose))):
+            
+            # print(file_path, root_path, save_dir)
+            
+            if root_path is not None:
+                file_path = file_path.replace(root_path+'/','')
+
+            save_path = os.path.join(save_dir, file_path)
+            
+            if save and audio_features is not None:
+                
+
+                #remove the root path from the file path
+                
+                
+                if 's3://' in save_dir:
+                    bucket, key = save_dir.replace("s3://", "").split("/", 1)
+                    key = f"{key}/{file_path}"
+                    
+                    # local_path = os.path.join(local_temp_dir, file_path)
+                    # os.makedirs(os.path.dirname(local_path), exist_ok=True)
+                    # np.save(local_path, audio_features.detach().cpu().numpy())
+                    
+                    pbar.set_description(f"Uploading features to s3://{bucket}/{key}") if verbose else None
+                    try:
+                        # client.upload_file(save_path, bucket, key)
+                        
+                        buffer = io.BytesIO()
+                        np.save(buffer, audio_features.detach().cpu().numpy())
+                        buffer.seek(0)
+                        client.put_object(Bucket=bucket, Key=key, Body=buffer)
+                    except Exception as e:
+                        print(f"Error uploading to s3: {e}") if verbose else None
+                        
+                    # os.remove(local_path)
+                else:
+                    pbar.set_description(f"Saving features in {save_path}, shape: {audio_features.shape}")
+                    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+                    if os.path.exists(save_path):
+                        os.remove(save_path)
+                    np.save(save_path, audio_features.detach().cpu().numpy())
+            
+            if not save and audio_features is not None:
+                pbar.set_description(f"{file_path}, shape: {audio_features.shape}")
+                pass
+                
+            audio_features_all.append(audio_features.detach().cpu()) if audio_features is not None else None
+            
+            counter += 1
+            if limit_n and counter >= limit_n:
+                break
+        
+        if skipped_count > 0:
+            print(f"Skipped {skipped_count} already processed items") if verbose else None
+        
+        try:   
+            print(f"Returning {len(audio_features_all)} features") if verbose else None
+            all_= torch.stack(audio_features_all)
+            print(f"Stacked features, shape: {all_.shape}") if verbose else None
+            return all_
+        
+        
+        except:
+            return None

steerable_retrieval/dataloading/loading_utils.py

@@ -0,0 +1,131 @@
+import torchaudio
+import soundfile as sf
+import numpy as np
+import torch
+import logging
+
+
+def _safe_sr_frames(path: str):
+    try:
+        info = sf.info(path)
+        return info.samplerate, info.frames
+    except sf.LibsndfileError:
+        # soundfile can't read some formats (e.g. mp3); fall back to torchaudio,
+        # which uses the same ffmpeg/soundfile backends without needing torchcodec.
+        info = torchaudio.info(path)
+        return info.sample_rate, info.num_frames
+
+
+def load_audio_chunk(path, target_n_samples, target_sr, start=None, verbose=False):
+    """
+    Load a chunk of audio from a file using torchaudio.
+    
+    Args:
+        path: Path to audio file
+        target_n_samples: Number of samples to load (at target_sr)
+        target_sr: Target sample rate
+        start: Starting frame (at original sr). If None, random start.
+        verbose: Print debug info
+    
+    Returns:
+        audio: Tensor of shape [n_samples, n_channels]
+    """
+    # Get audio metadata using soundfile
+    sr, frames = _safe_sr_frames(path)
+
+    print(f"length of audio in seconds: {frames/sr}") if verbose else None
+    print(f"Original sample rate: {sr}") if verbose else None
+
+    # Adjust for MP3 padding
+    if path.split(".")[-1].lower() == "mp3":
+        frames = frames - 8192
+
+    # Calculate how many frames to load at original sr
+    new_target_n_samples = int(target_n_samples * sr / target_sr)
+
+    print(f"New target n samples: {new_target_n_samples}") if verbose else None
+
+    # Random start if not specified
+    if start is None:
+        max_start = max(1, frames - new_target_n_samples)
+        start = np.random.randint(0, max_start)
+
+    # Load audio chunk with torchaudio
+    # torchaudio.load returns (waveform, sample_rate) where waveform is [channels, samples]
+    audio, loaded_sr = torchaudio.load(
+        path,
+        frame_offset=start,
+        num_frames=new_target_n_samples,
+        normalize=True,
+    )
+
+    # Resample if needed
+    if loaded_sr != target_sr:
+        audio = torchaudio.functional.resample(audio, loaded_sr, target_sr)
+        print(f"Resampled to {target_sr}, shape of audio: {audio.shape}") if verbose else None
+
+    # pad if needed
+    if audio.shape[1] < target_n_samples:
+        audio = torch.nn.functional.pad(audio, (0, target_n_samples - audio.shape[1]))
+
+    # Convert from [channels, samples] to [samples, channels]
+    audio = audio.T
+
+    return audio
+
+
+def load_full_audio(path, target_sr, verbose=False):
+    """
+    Load full audio file using torchaudio.
+    
+    Args:
+        path: Path to audio file
+        target_sr: Target sample rate
+        verbose: Print debug info
+    
+    Returns:
+        audio: Tensor of shape [n_channels, n_samples]
+    """
+    # Load with torchaudio - returns [channels, samples]
+    audio, sr = torchaudio.load(path, normalize=True)
+
+    # If stereo, average to mono
+    if audio.shape[0] == 2:
+        audio = audio.mean(dim=0, keepdim=True)
+
+
+    # Resample if needed
+    if sr != target_sr:
+        audio = torchaudio.functional.resample(audio, sr, target_sr)
+
+
+    return audio
+
+
+def load_full_and_split(path, target_sr, target_n_samples, hop=None, verbose=False):
+    """
+    Load full audio and split into overlapping chunks.
+    
+    Args:
+        path: Path to audio file
+        target_sr: Target sample rate
+        target_n_samples: Samples per chunk
+        hop: Hop size between chunks (default: target_n_samples)
+        verbose: Print debug info
+    
+    Returns:
+        audio: Tensor of shape [n_chunks, 1, target_n_samples]
+    """
+    hop = target_n_samples if hop is None else hop
+    audio = load_full_audio(path, target_sr, verbose=verbose)
+    audio = audio.squeeze()
+
+
+    # If audio is shorter than target, repeat
+    if audio.shape[0] < target_n_samples:
+        n_repeats = int(np.ceil(target_n_samples / audio.shape[0]))
+        audio = audio.repeat(n_repeats)
+
+    audio = audio.unfold(0, int(target_n_samples), int(hop)).unsqueeze(1)
+
+    return audio

steerable_retrieval/experiments/README.md

@@ -0,0 +1,35 @@
+# Offline Processing Experiments
+
+This directory contains offline experiment jobs that run outside training callbacks.
+
+## Stability (offline)
+
+Execution script:
+- `python steerable_retrieval/experiments/run_stability.py model_source.config_source=<WANDB_RUN_ID_OR_YAML_PATH> model_source.checkpoint_source=<CKPT> stability.checkpoint_paths=[<CKPT1>,<CKPT2>]`
+
+Launch as SageMaker processing job:
+- `python launch_stability.py model_source.config_source=<WANDB_RUN_ID_OR_YAML_PATH> model_source.checkpoint_source=<CKPT> stability.checkpoint_paths=[<CKPT1>,<CKPT2>]`
+
+## Concept Isolation (offline)
+
+Execution script:
+- `python steerable_retrieval/experiments/run_concept_isolation.py model_source.config_source=<WANDB_RUN_ID_OR_YAML_PATH> model_source.checkpoint_source=<CKPT> concepts_experiment=extract concepts.vocab_path=<VOCAB_TXT>`
+
+Score edit (add):
+- `python steerable_retrieval/experiments/run_concept_isolation.py model_source.config_source=<WANDB_RUN_ID_OR_YAML_PATH> model_source.checkpoint_source=<CKPT> concepts_experiment=score_edit_add concepts.vocab_path=<VOCAB_TXT>`
+
+Score edit (suppress):
+- `python steerable_retrieval/experiments/run_concept_isolation.py model_source.config_source=<WANDB_RUN_ID_OR_YAML_PATH> model_source.checkpoint_source=<CKPT> concepts_experiment=score_edit_suppress concepts.vocab_path=<VOCAB_TXT>`
+
+Launch as SageMaker processing job:
+- `python launch_concept_isolation.py model_source.config_source=<WANDB_RUN_ID_OR_YAML_PATH> model_source.checkpoint_source=<CKPT> concepts_experiment=extract concepts.vocab_path=<VOCAB_TXT>`
+
+## Notes
+
+- Both jobs save run config and outputs under `output_dir` (default `${paths.output_dir}`).
+- Stability logging to W&B is controlled by `wandb.*` keys in `configs/experiment/stability.yaml`.
+- Concept isolation is now a dispatcher with `concepts_experiment=extract|score_edit_add|score_edit_suppress`.
+- `model_source.config_source` can be a W&B run id, a local `.yaml` file path, or an `s3://...yaml` path.
+- If `config_source` contains `.yaml`, the model config is read from that YAML; otherwise it is fetched from W&B.
+- `model_source.checkpoint_source` is the checkpoint path loaded directly (local path or S3 path).
+

steerable_retrieval/experiments/__init__.py

@@ -0,0 +1,2 @@
+"""Offline experiment entrypoints and utilities."""
+

steerable_retrieval/experiments/common.py

@@ -0,0 +1,343 @@
+from __future__ import annotations
+
+import importlib
+import logging
+import os
+from pathlib import Path
+from urllib.parse import urlparse
+from typing import Any, Callable, Dict, List, Optional
+
+import hydra
+import torch
+from omegaconf import DictConfig, OmegaConf
+
+from steerable_retrieval.callbacks.energy import load_state_dict_any
+from steerable_retrieval.utils.pylogger import RankedLogger
+
+log = RankedLogger(__name__, rank_zero_only=True)
+
+
+def _silence_huggingface_loading() -> None:
+    """Reduce HuggingFace/Transformers load-time noise (progress bars + logs)."""
+    os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
+    os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
+    logging.getLogger("transformers").setLevel(logging.ERROR)
+    logging.getLogger("huggingface_hub").setLevel(logging.ERROR)
+    try:
+        from transformers.utils import logging as hf_logging
+
+        hf_logging.set_verbosity_error()
+    except Exception:
+        pass
+
+
+def ensure_dir(path: str) -> str:
+    os.makedirs(path, exist_ok=True)
+    return path
+
+
+def save_config(cfg: DictConfig, output_dir: str, filename: str = "config.yaml") -> str:
+    ensure_dir(output_dir)
+    out_path = os.path.join(output_dir, filename)
+    OmegaConf.save(cfg, out_path, resolve=True)
+    return out_path
+
+
+def resolve_checkpoint_path(model_source_cfg: DictConfig) -> str:
+    ckpt = model_source_cfg.get("checkpoint_source")
+    if ckpt is None:
+        # Backward compatibility with older experiment configs.
+        ckpt = model_source_cfg.get("checkpoint_path")
+    if ckpt:
+        return str(ckpt)
+
+    run_id = model_source_cfg.get("run_id")
+    if run_id is None:
+        run_id = model_source_cfg.get("config_source")
+    if run_id is None:
+        run_id = model_source_cfg.get("wandb_run_id")
+    if not run_id:
+        raise ValueError(
+            "Missing model_source.run_id (or compatible fallback). "
+            "Set model_source.checkpoint_source explicitly or provide a run id for auto checkpoint resolution."
+        )
+    run_id = str(run_id)
+    if ".yaml" in run_id:
+        raise ValueError(
+            "Auto checkpoint resolution expects model_source.config_source to be a run id, "
+            "not a YAML path. Set model_source.checkpoint_source explicitly."
+        )
+
+    s3_prefix = str(model_source_cfg.get("checkpoint_s3_prefix", "s3://your-bucket/checkpoints")).rstrip("/")
+    parsed = urlparse(s3_prefix)
+    if parsed.scheme != "s3" or not parsed.netloc:
+        raise ValueError(
+            f"Invalid model_source.checkpoint_s3_prefix '{s3_prefix}'. Expected an s3:// URI."
+        )
+    bucket = parsed.netloc
+    prefix = parsed.path.lstrip("/")
+    key_prefix = f"{prefix}/{run_id}/checkpoints/".lstrip("/")
+
+    try:
+        import boto3
+    except Exception as exc:
+        raise RuntimeError("boto3 is required to auto-resolve checkpoint path from S3.") from exc
+
+    client = boto3.client("s3")
+    paginator = client.get_paginator("list_objects_v2")
+    candidates = []
+    for page in paginator.paginate(Bucket=bucket, Prefix=key_prefix):
+        for obj in page.get("Contents", []):
+            key = str(obj.get("Key", ""))
+            filename = key.rsplit("/", 1)[-1]
+            if "epoch" in filename and (key.endswith(".pt") or key.endswith(".ckpt")):
+                candidates.append(obj)
+
+    if not candidates:
+        raise FileNotFoundError(
+            "No checkpoint file found under "
+            f"s3://{bucket}/{key_prefix} matching '*epoch*.pt' or '*epoch*.ckpt'."
+        )
+
+    latest = max(candidates, key=lambda x: x.get("LastModified"))
+    latest_key = str(latest["Key"])
+    resolved_ckpt = f"s3://{bucket}/{latest_key}"
+    log.info(f"Auto-resolved checkpoint from run id '{run_id}': {resolved_ckpt}")
+    return resolved_ckpt
+
+
+def resolve_run_id(run_id: str, entity: Optional[str], project: Optional[str]) -> str:
+    if not entity or not project:
+        raise ValueError(
+            "wandb.entity and wandb.project are required to resolve run id via "
+            "config.sagemaker_job_name."
+        )
+    try:
+        import wandb
+    except Exception as exc:
+        raise RuntimeError("wandb is required to resolve model_source.config_source.") from exc
+
+    api = wandb.Api()
+    runs = api.runs(
+        path=f"{entity}/{project}",
+        filters={"config.sagemaker_job_name": run_id},
+        per_page=1,
+    )
+    runs = list(runs)
+    if not runs:
+        raise ValueError(
+            f"No W&B run found in {entity}/{project} with config.sagemaker_job_name='{run_id}'."
+        )
+    resolved = str(runs[0].id)
+    log.info(
+        f"Resolved model_source.config_source '{run_id}' to W&B run id '{resolved}' "
+        f"via config.sagemaker_job_name."
+    )
+    return resolved
+
+
+def _fetch_model_cfg_from_wandb(run_id: str, entity: Optional[str], project: Optional[str]) -> DictConfig:
+    resolved_run_id = resolve_run_id(run_id=run_id, entity=entity, project=project)
+    run_path = f"{entity}/{project}/{resolved_run_id}"
+    import wandb
+    api = wandb.Api()
+    run = api.run(run_path)
+    run_cfg = OmegaConf.create(run.config)
+    model_cfg = run_cfg.get("model")
+    if model_cfg is None:
+        raise KeyError(f"No 'model' key found in wandb run config for {run_path}.")
+    model_cfg = OmegaConf.create(model_cfg)
+    try:
+        OmegaConf.resolve(model_cfg)
+    except Exception:
+        log.warning(
+            "Could not fully resolve interpolations inside wandb run model config. "
+            "Proceeding with unresolved values."
+        )
+    return OmegaConf.create(OmegaConf.to_container(model_cfg, resolve=False))
+
+
+def _fetch_yaml_cfg_from_s3(s3_uri: str) -> DictConfig:
+    try:
+        import boto3
+    except Exception as exc:
+        raise RuntimeError("boto3 is required to read YAML model config from S3.") from exc
+
+    parsed = urlparse(s3_uri)
+    bucket = parsed.netloc
+    key = parsed.path.lstrip("/")
+    if not bucket or not key:
+        raise ValueError(f"Invalid S3 URI for model_source.config_source: {s3_uri}")
+
+    client = boto3.client("s3")
+    body = client.get_object(Bucket=bucket, Key=key)["Body"].read().decode("utf-8")
+    cfg = OmegaConf.create(body)
+    OmegaConf.resolve(cfg)
+    return cfg
+
+
+def _fetch_model_cfg_from_yaml(config_source: str) -> DictConfig:
+    if config_source.startswith("s3://"):
+        cfg = _fetch_yaml_cfg_from_s3(config_source)
+    else:
+        cfg = OmegaConf.load(config_source)
+        OmegaConf.resolve(cfg)
+
+    model_cfg = cfg.get("model")
+    # Accept either a full experiment config (with "model") or a model-only config.
+    if model_cfg is None:
+        model_cfg = cfg
+    return OmegaConf.create(OmegaConf.to_container(model_cfg, resolve=True))
+
+
+def resolve_model_from_source(cfg: DictConfig) -> str:
+    model_source_cfg = cfg.get("model_source")
+    if model_source_cfg is None:
+        raise ValueError("Missing cfg.model_source.")
+
+    run_id = model_source_cfg.get("run_id")
+    config_source = model_source_cfg.get("config_source")
+    if config_source is None:
+        # Backward compatibility with older experiment configs.
+        config_source = model_source_cfg.get("wandb_run_id")
+    if not run_id and not config_source:
+        raise ValueError(
+            "model_source.run_id must be set (preferred), or set model_source.config_source "
+            "to a W&B id or a .yaml path (local/S3)."
+        )
+
+    if run_id:
+        run_id = str(run_id)
+        model_cfg = _fetch_model_cfg_from_wandb(
+            run_id=run_id,
+            entity=cfg.get("wandb", {}).get("entity"),
+            project=cfg.get("wandb", {}).get("project"),
+        )
+        resolved_source = run_id
+    else:
+        config_source = str(config_source)
+        if ".yaml" in config_source:
+            model_cfg = _fetch_model_cfg_from_yaml(config_source)
+        else:
+            model_cfg = _fetch_model_cfg_from_wandb(
+                run_id=config_source,
+                entity=cfg.get("wandb", {}).get("entity"),
+                project=cfg.get("wandb", {}).get("project"),
+            )
+        resolved_source = config_source
+
+    # Apply optional local/runtime overrides after fetching model config from source.
+    # This is useful for local CPU smoke tests when source configs are CUDA-oriented.
+    post_overrides = cfg.get("model_post_overrides", {})
+    if not torch.cuda.is_available():
+        model_cfg.sae_encoder.device = 'cpu'
+        model_cfg.sae_decoder.device = 'cpu'
+        model_cfg.audio_encoder.device = 'cpu'
+        model_cfg.text_encoder.device = 'cpu'
+    if post_overrides:
+        model_cfg = OmegaConf.merge(model_cfg, post_overrides)
+
+    cfg.model = model_cfg
+    return resolved_source
+
+
+def instantiate_model_and_load(cfg: DictConfig, device: torch.device):
+    _silence_huggingface_loading()
+    model = hydra.utils.instantiate(cfg.model)
+    ckpt_path = resolve_checkpoint_path(cfg.model_source)
+    state_dict = load_state_dict_any(ckpt_path, map_location="cpu")
+    model_state_keys = set(model.state_dict().keys())
+    checkpoint_keys = set(state_dict.keys()) if isinstance(state_dict, dict) else set()
+    matched_keys = model_state_keys.intersection(checkpoint_keys)
+    if not matched_keys:
+        raise RuntimeError(
+            "Checkpoint appears incompatible: zero matching keys between model and checkpoint state_dict."
+        )
+    matched_key_names = sorted(matched_keys)
+    load_result = model.load_state_dict(state_dict, strict=False)
+    load_report = {
+        "checkpoint_keys": len(checkpoint_keys),
+        "model_keys": len(model_state_keys),
+        "matched_keys": len(matched_keys),
+        "matched_key_names": matched_key_names,
+        "missing_keys": len(getattr(load_result, "missing_keys", [])),
+        "unexpected_keys": len(getattr(load_result, "unexpected_keys", [])),
+    }
+    setattr(model, "_checkpoint_load_report", load_report)
+    model = model.to(device)
+    model.eval()
+    preview_limit = 30
+    matched_preview = matched_key_names[:preview_limit]
+    remaining = len(matched_key_names) - len(matched_preview)
+    preview_suffix = f" ... (+{remaining} more)" if remaining > 0 else ""
+    log.info(
+        f"Loaded model from checkpoint: {ckpt_path} | "
+        f"matched={load_report['matched_keys']}/{load_report['model_keys']} "
+        f"missing={load_report['missing_keys']} unexpected={load_report['unexpected_keys']} "
+        f"| matched_keys={matched_preview}{preview_suffix}"
+    )
+    return model, ckpt_path
+
+
+def read_vocab_lines(vocab_path: str) -> List[str]:
+    with open(vocab_path, "r", encoding="utf-8") as f:
+        concepts = [line.strip() for line in f if line.strip()]
+    if not concepts:
+        raise ValueError(f"No concepts found in vocabulary file: {vocab_path}")
+    return concepts
+
+
+def resolve_callable(path_str: str) -> Callable[..., Any]:
+    if ":" in path_str:
+        module_name, fn_name = path_str.split(":", 1)
+    else:
+        module_name, fn_name = path_str.rsplit(".", 1)
+    module = importlib.import_module(module_name)
+    fn = getattr(module, fn_name)
+    if not callable(fn):
+        raise TypeError(f"Resolved object is not callable: {path_str}")
+    return fn
+
+
+def maybe_start_wandb(cfg: DictConfig, job_type: str, name: str, output_dir: str):
+    wb_cfg = cfg.get("wandb", {})
+    if not wb_cfg.get("enabled", True):
+        return None
+    try:
+        import wandb
+    except Exception:
+        log.warning("wandb not available, skipping remote logging.")
+        return None
+
+    run = wandb.init(
+        project=wb_cfg.get("project", "SpaMR"),
+        entity=wb_cfg.get("entity"),
+        name=name,
+        group=wb_cfg.get("group"),
+        job_type=job_type,
+        tags=list(wb_cfg.get("tags", [])),
+        dir=output_dir,
+        config=OmegaConf.to_container(cfg, resolve=True),
+    )
+    return run
+
+
+def finish_wandb(run) -> None:
+    if run is None:
+        return
+    try:
+        import wandb
+
+        wandb.finish()
+    except Exception:
+        return
+
+
+def save_json(path: str, payload: Dict[str, Any]) -> None:
+    import json
+
+    ensure_dir(str(Path(path).parent))
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(payload, f, indent=2, sort_keys=True)
+
+

steerable_retrieval/experiments/concepts/__init__.py

@@ -0,0 +1,2 @@
+"""Concept isolation experiment implementations."""
+

steerable_retrieval/experiments/concepts/extract.py

@@ -0,0 +1,899 @@
+from __future__ import annotations
+
+import importlib
+import math
+import os
+import re
+from functools import partial
+from typing import Any, Dict
+
+import torch
+from omegaconf import DictConfig, OmegaConf
+import torch.nn.functional as F
+
+from steerable_retrieval.callbacks.energy import BridgeScoreCallback, get_dictionary_from_lightningsae, load_state_dict_any
+from steerable_retrieval.experiments.common import ensure_dir, read_vocab_lines
+from steerable_retrieval.utils import RankedLogger
+
+log = RankedLogger(__name__, rank_zero_only=True)
+
+
+def _slugify(value: str) -> str:
+    slug = re.sub(r"[^A-Za-z0-9._-]+", "-", str(value)).strip("-")
+    return slug or "na"
+
+
+def _flatten_params(prefix: str, value: Any, items: list[tuple[str, str]]) -> None:
+    if isinstance(value, dict):
+        for key in sorted(value.keys()):
+            child_prefix = f"{prefix}.{key}" if prefix else str(key)
+            _flatten_params(child_prefix, value[key], items)
+        return
+    items.append((prefix or "param", str(value)))
+
+
+def _auto_experiment_id(cfg: DictConfig) -> str:
+    extractor_cfg = cfg.concepts_experiment.get("extractor")
+    params_cfg = extractor_cfg.get("params", {}) if extractor_cfg is not None else {}
+    params = OmegaConf.to_container(params_cfg, resolve=True) if params_cfg is not None else {}
+    if not isinstance(params, dict):
+        params = {"value": params}
+
+    method = params.get("selection_method")
+    if method is None:
+        target = str(extractor_cfg.get("target", "extract")) if extractor_cfg is not None else "extract"
+        method = target.rsplit(".", 1)[-1]
+    method_slug = _slugify(method)
+
+    flat_items: list[tuple[str, str]] = []
+    _flatten_params("", params, flat_items)
+    parts = [f"{_slugify(k)}-{_slugify(v)}" for k, v in flat_items if str(v) != ""]
+    suffix = "__".join(parts)
+    out = f"{method_slug}__{suffix}" if suffix else method_slug
+    return out[:180]
+
+
+def _resolve_output_dir(cfg: DictConfig, output_dir: str, source_run_id: str) -> tuple[str, str, str]:
+    run_id = cfg.get("model_source", {}).get("run_id") or source_run_id or "unknown_run"
+    run_slug = _slugify(run_id)
+    experiment_id_cfg = cfg.get("experiment_id")
+    experiment_id = _slugify(experiment_id_cfg) if experiment_id_cfg else _auto_experiment_id(cfg)
+    full_output_dir = ensure_dir(os.path.join(output_dir, run_slug, experiment_id))
+    return full_output_dir, run_slug, experiment_id
+
+
+def _sort_unique_by_abs_desc(indices: torch.Tensor, values: torch.Tensor) -> torch.Tensor:
+    if indices.numel() == 0:
+        return indices.to(dtype=torch.long)
+    uniq = torch.unique(indices.to(dtype=torch.long))
+    abs_vals = values[uniq].abs()
+    order = torch.argsort(abs_vals, descending=True)
+    return uniq[order]
+
+
+def _cap_indices(indices: torch.Tensor, max_selected: int | None) -> torch.Tensor:
+    if max_selected is None:
+        return indices
+    if max_selected <= 0:
+        return indices[:0]
+    return indices[: int(max_selected)]
+
+
+def get_selected_indices_abs_mass(
+    cosine_similarities: torch.Tensor,
+    *,
+    eta: float = 0.95,
+    max_selected: int | None = None,
+):
+    abs_cosine = cosine_similarities.abs()
+    selected_indices = torch.empty(0, dtype=torch.long, device=cosine_similarities.device)
+    k95 = 0
+    if abs_cosine.numel() > 0:
+        sorted_abs_values, sorted_abs_indices = torch.sort(abs_cosine, descending=True)
+        total_abs_mass = sorted_abs_values.sum()
+        if float(total_abs_mass.item()) > 0.0:
+            cumsum_abs = torch.cumsum(sorted_abs_values, dim=0)
+            threshold = eta * total_abs_mass
+            cutoff_candidates = torch.nonzero(cumsum_abs >= threshold, as_tuple=False)
+            cutoff = int(cutoff_candidates[0].item()) + 1 if cutoff_candidates.numel() > 0 else sorted_abs_values.numel()
+        else:
+            cutoff = 0
+        selected_indices = sorted_abs_indices[:cutoff]
+        k95 = cutoff
+    selected_indices = _cap_indices(selected_indices, max_selected=max_selected)
+    return selected_indices, {
+        "method": "abs_mass",
+        "k95": int(k95),
+        "eta": float(eta),
+        "num_selected": int(selected_indices.numel()),
+    }
+
+
+def get_selected_indices_zscore(
+    cosine_similarities: torch.Tensor,
+    *,
+    tau: float = 2.5,
+    max_selected: int | None = None,
+):
+    x = cosine_similarities
+    if x.numel() == 0:
+        return x.new_empty(0, dtype=torch.long), {"method": "zscore", "tau": float(tau), "num_selected": 0}
+    mu = x.mean()
+    sigma = x.std(unbiased=False)
+    sigma_val = float(sigma.item())
+    if sigma_val <= 0.0:
+        return x.new_empty(0, dtype=torch.long), {
+            "method": "zscore",
+            "tau": float(tau),
+            "mu": float(mu.item()),
+            "sigma": sigma_val,
+            "num_selected": 0,
+            "num_pos": 0,
+        }
+
+    upper = mu + tau * sigma
+    lower = mu - tau * sigma
+    pos_idx = torch.nonzero(x >= upper, as_tuple=False).flatten()
+    # if zero, select the max positive value
+    if pos_idx.numel() == 0:
+        pos_idx = torch.argmax(x)
+    selected = _sort_unique_by_abs_desc(pos_idx, x)
+    selected = _cap_indices(selected, max_selected=max_selected)
+    return selected, {
+        "method": "zscore",
+        "tau": float(tau),
+        "mu": float(mu.item()),
+        "sigma": sigma_val,
+        "threshold_upper": float(upper.item()),
+        "threshold_lower": float(lower.item()),
+        "num_pos": int(pos_idx.numel()),
+        "num_selected": int(selected.numel()),
+    }
+
+
+def get_selected_indices_fdr(
+    cosine_similarities: torch.Tensor,
+    *,
+    q0: float = 0.05,
+    max_selected: int | None = None,
+):
+    x = cosine_similarities
+    if x.numel() == 0:
+        return x.new_empty(0, dtype=torch.long), {"method": "fdr", "q0": float(q0), "num_selected": 0}
+
+    mu = x.mean()
+    sigma = x.std(unbiased=False).clamp_min(1e-12)
+    z = (x - mu) / sigma
+    # Two-sided p-values under normal null.
+    pvals = torch.erfc(z.abs() / math.sqrt(2.0)).clamp_min(1e-12).clamp_max(1.0)
+    m = int(pvals.numel())
+    sorted_pvals, sorted_idx = torch.sort(pvals, descending=False)
+    ranks = torch.arange(1, m + 1, device=x.device, dtype=x.dtype)
+    thresholds = (q0 * ranks) / float(m)
+    passed = sorted_pvals <= thresholds
+    if not bool(passed.any()):
+        selected = x.new_empty(0, dtype=torch.long)
+        p_cut = None
+    else:
+        k = int(torch.nonzero(passed, as_tuple=False)[-1].item()) + 1
+        p_cut = float(sorted_pvals[k - 1].item())
+        selected = sorted_idx[:k]
+
+    selected = _sort_unique_by_abs_desc(selected, x)
+    selected = _cap_indices(selected, max_selected=max_selected)
+    return selected, {
+        "method": "fdr",
+        "q0": float(q0),
+        "mu": float(mu.item()),
+        "sigma": float(sigma.item()),
+        "p_cut": p_cut,
+        "num_selected": int(selected.numel()),
+        "num_tests": m,
+    }
+
+
+def get_selected_indices_gmm(
+    cosine_similarities: torch.Tensor,
+    *,
+    posterior_threshold: float = 0.8,
+    max_selected: int | None = None,
+    max_iter: int = 50,
+):
+    x = cosine_similarities
+    if x.numel() == 0:
+        return x.new_empty(0, dtype=torch.long), {"method": "gmm", "num_selected": 0}
+    x1 = x.reshape(-1, 1)
+    n = x1.shape[0]
+    mu = torch.quantile(x, torch.tensor([0.25, 0.75], device=x.device, dtype=x.dtype)).reshape(2, 1)
+    var0 = x.var(unbiased=False).clamp_min(1e-6)
+    var = torch.full((2, 1), var0, device=x.device, dtype=x.dtype)
+    pi = torch.full((2, 1), 0.5, device=x.device, dtype=x.dtype)
+    two_pi = torch.tensor(2.0 * math.pi, device=x.device, dtype=x.dtype)
+
+    for _ in range(max_iter):
+        # E-step
+        norm = (1.0 / torch.sqrt(two_pi * var)) * torch.exp(-0.5 * ((x1 - mu) ** 2) / var)
+        weighted = pi * norm
+        denom = weighted.sum(dim=0, keepdim=True).clamp_min(1e-12)
+        resp = weighted / denom
+        # M-step
+        Nk = resp.sum(dim=1, keepdim=True).clamp_min(1e-12)
+        pi = Nk / float(n)
+        mu = (resp @ x1) / Nk
+        centered = x1.unsqueeze(0) - mu.unsqueeze(1)
+        var = ((resp.unsqueeze(2) * (centered**2)).sum(dim=1) / Nk).clamp_min(1e-6)
+
+    # "Relevant" component: larger absolute mean.
+    relevant_idx = int(torch.argmax(mu.abs()).item())
+    posterior_relevant = resp[relevant_idx]
+    selected = torch.nonzero(posterior_relevant > posterior_threshold, as_tuple=False).flatten()
+    selected = _sort_unique_by_abs_desc(selected, x)
+    selected = _cap_indices(selected, max_selected=max_selected)
+    return selected, {
+        "method": "gmm",
+        "posterior_threshold": float(posterior_threshold),
+        "means": [float(mu[0].item()), float(mu[1].item())],
+        "stds": [float(torch.sqrt(var[0]).item()), float(torch.sqrt(var[1]).item())],
+        "weights": [float(pi[0].item()), float(pi[1].item())],
+        "relevant_component": relevant_idx,
+        "num_selected": int(selected.numel()),
+    }
+
+
+def get_selected_indices_sparse_pursuit(
+    cosine_similarities: torch.Tensor,
+    *,
+    target_k: int | None = None,
+    l1_lambda: float = 0.0,
+    max_selected: int | None = None,
+):
+    x = cosine_similarities
+    if x.numel() == 0:
+        return x.new_empty(0, dtype=torch.long), {"method": "sparse_pursuit", "num_selected": 0}
+
+    if target_k is not None and target_k > 0:
+        k = min(int(target_k), int(x.numel()))
+        selected = torch.topk(x.abs(), k=k).indices
+        mode = "topk_abs"
+        threshold = float(torch.topk(x.abs(), k=k).values[-1].item()) if k > 0 else 0.0
+    else:
+        lam = float(l1_lambda)
+        selected = torch.nonzero(x.abs() >= lam, as_tuple=False).flatten()
+        mode = "l1_threshold"
+        threshold = lam
+
+    selected = _sort_unique_by_abs_desc(selected, x)
+    selected = _cap_indices(selected, max_selected=max_selected)
+    return selected, {
+        "method": "sparse_pursuit",
+        "mode": mode,
+        "threshold": threshold,
+        "target_k": None if target_k is None else int(target_k),
+        "l1_lambda": float(l1_lambda),
+        "num_selected": int(selected.numel()),
+    }
+
+
+def get_selected_indices_quantile(
+    cosine_similarities: torch.Tensor,
+    *,
+    q: float = 0.05,
+    positive_only: bool = True,
+    max_selected: int | None = None,
+):
+    x = cosine_similarities
+    if x.numel() == 0:
+        return x.new_empty(0, dtype=torch.long), {"method": "quantile", "q": float(q), "num_selected": 0}
+
+    q = float(q)
+    if not (0.0 < q < 1.0):
+        raise ValueError(f"quantile selection expects q in (0, 1), got {q}")
+
+    if positive_only:
+        pool = x[x > 0]
+        if pool.numel() == 0:
+            return x.new_empty(0, dtype=torch.long), {
+                "method": "quantile",
+                "q": q,
+                "positive_only": True,
+                "threshold": None,
+                "num_selected": 0,
+            }
+        threshold = torch.quantile(pool, 1.0 - q)
+        selected = torch.nonzero(x >= threshold, as_tuple=False).flatten()
+        selected = selected[x[selected] > 0]
+    else:
+        threshold = torch.quantile(x.abs(), 1.0 - q)
+        selected = torch.nonzero(x.abs() >= threshold, as_tuple=False).flatten()
+
+    selected = _sort_unique_by_abs_desc(selected, x)
+    selected = _cap_indices(selected, max_selected=max_selected)
+    return selected, {
+        "method": "quantile",
+        "q": q,
+        "positive_only": bool(positive_only),
+        "threshold": float(threshold.item()) if threshold is not None else None,
+        "num_selected": int(selected.numel()),
+    }
+
+
+def get_selected_indices(
+    cosine_similarities: torch.Tensor,
+    *,
+    method: str = "abs_mass",
+    **kwargs,
+):
+    method_key = str(method).lower()
+    if method_key in {"abs_mass", "cosine_mass", "k95"}:
+        return get_selected_indices_abs_mass(cosine_similarities, **kwargs)
+    if method_key in {"zscore", "z_score"}:
+        return get_selected_indices_zscore(cosine_similarities, **kwargs)
+    if method_key in {"fdr", "bh"}:
+        return get_selected_indices_fdr(cosine_similarities, **kwargs)
+    if method_key in {"gmm", "mixture"}:
+        return get_selected_indices_gmm(cosine_similarities, **kwargs)
+    if method_key in {"sparse_pursuit", "lasso", "topk"}:
+        return get_selected_indices_sparse_pursuit(cosine_similarities, **kwargs)
+    if method_key in {"quantile", "q"}:
+        return get_selected_indices_quantile(cosine_similarities, **kwargs)
+    raise ValueError(f"Unknown selection method '{method}'.")
+
+
+def attention(q, k, v):
+    return torch.nn.functional.softmax(q @ k.T / math.sqrt(q.shape[-1]), dim=-1) @ v
+
+
+def _attention_over_neurons(
+    query: torch.Tensor,
+    keys: torch.Tensor,
+    tau: float = 0.07,
+    normalize: bool = True,
+) -> torch.Tensor:
+    """
+    Attention over neuron prototypes.
+
+    Args
+    ----
+    query: [1, d] or [d]
+    keys:  [K, d]
+    tau: temperature (smaller => peakier)
+    normalize: if True, use cosine logits via L2-normalization
+
+    Returns
+    -------
+    alpha: [K] attention weights over neurons
+    """
+    if query.dim() == 1:
+        query = query.unsqueeze(0)  # [1, d]
+    assert query.dim() == 2 and keys.dim() == 2, (query.shape, keys.shape)
+    assert query.shape[-1] == keys.shape[-1], (query.shape, keys.shape)
+
+    if normalize:
+        q = F.normalize(query, dim=-1)          # [1, d]
+        k = F.normalize(keys, dim=-1)           # [K, d]
+        logits = (q @ k.t()).squeeze(0)         # [K], cosine logits
+    else:
+        # scaled dot-product in raw space
+        logits = (query @ keys.t()).squeeze(0) / math.sqrt(query.shape[-1])  # [K]
+
+    alpha = F.softmax(logits / tau, dim=-1)     # [K]
+    return alpha
+
+
+def _raw_cosine_for_concept(
+    *,
+    model,
+    concept: str,
+    device: torch.device,
+    text_encoder=None,
+    basis_chunk_size: int = 256,
+    attention: bool = False,
+    attention_tau: float = 0.07,
+    attention_normalize: bool = True,
+    attention_score_mode: str = "alpha",
+):
+    """Compute raw cosine-like scores for one concept with optional attention reweighting."""
+    encoder = text_encoder if text_encoder is not None else getattr(model, "text_encoder", None)
+    if encoder is None:
+        raise ValueError("A text encoder must be provided when model.text_encoder is not available.")
+
+    text_embedding = encoder([concept])  # [1, d]
+    sae_text_out = model.sae_encoder(text_embedding)
+    if isinstance(sae_text_out, (tuple, list)):
+        sae_text_activations = sae_text_out[1] if len(sae_text_out) > 1 else sae_text_out[0]
+    else:
+        sae_text_activations = sae_text_out
+
+    dict_size = int(model.sae_decoder.W_dec.shape[0])
+    if basis_chunk_size <= 0:
+        basis_chunk_size = dict_size
+
+    basis = []
+    text_vec = text_embedding
+    for start in range(0, dict_size, basis_chunk_size):
+        end = min(start + basis_chunk_size, dict_size)
+        chunk = end - start
+        eye_chunk = torch.zeros((chunk, dict_size), device=device, dtype=text_vec.dtype)
+        eye_chunk[torch.arange(chunk, device=device), torch.arange(start, end, device=device)] = 1
+        basis_chunk = model.sae_decoder(eye_chunk)
+        basis.append(basis_chunk)
+    basis = torch.cat(basis, dim=0)  # [K, d]
+
+    if attention:
+        alpha = _attention_over_neurons(
+            query=text_vec, keys=basis, tau=attention_tau, normalize=attention_normalize
+        )
+        cosine_similarities = F.cosine_similarity(text_vec.expand_as(basis), basis, dim=-1)
+        if attention_score_mode == "alpha":
+            cosine_similarities = alpha
+        elif attention_score_mode == "alpha_cos":
+            cosine_similarities = alpha * cosine_similarities
+        elif attention_score_mode == "cos":
+            pass
+        else:
+            raise ValueError(
+                f"Unknown attention_score_mode={attention_score_mode}. "
+                "Use one of {'alpha','alpha_cos','cos'}."
+            )
+    else:
+        cosine_similarities = F.cosine_similarity(text_vec.expand_as(basis), basis, dim=-1)
+
+    return cosine_similarities, sae_text_activations
+
+
+def _cosine_probe_vector_and_stats(
+    *,
+    model,
+    concept: str,
+    concept_index: int,
+    device: torch.device,
+    selection_method: str = "abs_mass",
+    selection_kwargs: Dict[str, Any] | None = None,
+    basis_chunk_size: int = 256,
+    attention: bool = False,
+    attention_tau: float = 0.07,
+    attention_normalize: bool = True,
+    attention_score_mode: str = "alpha",  # {"alpha", "alpha_cos", "cos"}
+):
+    del concept_index  # Kept for extractor signature consistency.
+    cosine_similarities, sae_text_activations = _raw_cosine_for_concept(
+        model=model,
+        concept=concept,
+        device=device,
+        basis_chunk_size=basis_chunk_size,
+        attention=attention,
+        attention_tau=attention_tau,
+        attention_normalize=attention_normalize,
+        attention_score_mode=attention_score_mode,
+    )
+
+    # --- Selection (unchanged) ---
+    selection_kwargs = dict(selection_kwargs or {})
+
+    selected_indices, selection_stats = get_selected_indices(
+        cosine_similarities,
+        method=selection_method,
+        **selection_kwargs,
+    )
+
+    mask = torch.zeros_like(cosine_similarities)
+    if selected_indices.numel() > 0:
+        mask[selected_indices] = 1
+    masked_cosine = cosine_similarities * mask
+
+    probs = torch.softmax(cosine_similarities.abs(), dim=0)
+    entropy = float((-(probs * torch.log(probs.clamp_min(1e-12))).sum()).item())
+    stats = dict(selection_stats)
+    stats["entropy"] = entropy
+
+    return {
+        "raw_cosine_similarities": cosine_similarities,
+        "masked_cosine": masked_cosine,
+        "sae_text_activations": sae_text_activations,
+        "stats": stats,
+    }
+
+
+@torch.no_grad()
+def cosine_with_tfidf(
+    *,
+    model,
+    concepts: list[str],
+    device: torch.device,
+    tau: float = 4.0,
+    basis_chunk_size: int = 256,
+    attention: bool = False,
+    attention_tau: float = 0.07,
+    attention_normalize: bool = True,
+    attention_score_mode: str = "alpha",
+):
+    """
+    Two-pass cosine extraction with global TF-IDF neuron reweighting.
+
+    Steps:
+      1) raw cosine per concept
+      2) z-score mask on raw cosine
+      3) compute global neuron TF-IDF weights from first-pass masks
+      4) apply weights to raw cosine
+      5) z-score mask again on weighted cosine
+    """
+    if not concepts:
+        return {}
+
+    raw_by_concept: Dict[str, torch.Tensor] = {}
+    sae_by_concept: Dict[str, torch.Tensor] = {}
+    first_mask_by_concept: Dict[str, torch.Tensor] = {}
+
+    # Pass 1: raw cosine + first z-score mask
+    for concept in concepts:
+        raw_cos, sae_text_acts = _raw_cosine_for_concept(
+            model=model,
+            concept=concept,
+            device=device,
+            basis_chunk_size=basis_chunk_size,
+            attention=attention,
+            attention_tau=attention_tau,
+            attention_normalize=attention_normalize,
+            attention_score_mode=attention_score_mode,
+        )
+        idx1, stats1 = get_selected_indices_zscore(raw_cos, tau=tau)
+        mask1 = torch.zeros_like(raw_cos)
+        if idx1.numel() > 0:
+            mask1[idx1] = 1
+        raw_by_concept[concept] = raw_cos
+        sae_by_concept[concept] = sae_text_acts
+        first_mask_by_concept[concept] = mask1
+
+    # Global TF-IDF weights over neurons from first-pass masks.
+    first_mask_matrix = torch.stack([first_mask_by_concept[c] for c in concepts], dim=0)  # [C, K]
+    raw_matrix = torch.stack([raw_by_concept[c] for c in concepts], dim=0)  # [C, K]
+    n_concepts = first_mask_matrix.shape[0]
+    doc_freq = first_mask_matrix.sum(dim=0)  # [K]
+    tf = (raw_matrix.abs() * first_mask_matrix).sum(dim=0) / max(n_concepts, 1)  # [K]
+    idf = torch.log((n_concepts + 1.0) / (doc_freq + 1.0)) + 1.0  # [K]
+    neuron_weights = tf * idf
+    # Normalize to keep scales comparable across runs.
+    neuron_weights = neuron_weights / neuron_weights.mean().clamp_min(1e-12)
+
+    out = {}
+    # Pass 2: reweight raw cosine + second z-score mask
+    for concept in concepts:
+        raw_cos = raw_by_concept[concept]
+        weighted_cos = raw_cos * neuron_weights
+        idx2, stats2 = get_selected_indices_zscore(weighted_cos, tau=tau)
+        mask2 = torch.zeros_like(weighted_cos)
+        if idx2.numel() > 0:
+            mask2[idx2] = 1
+        masked_weighted = weighted_cos * mask2
+
+        probs = torch.softmax(weighted_cos.abs(), dim=0)
+        entropy = float((-(probs * torch.log(probs.clamp_min(1e-12))).sum()).item())
+        stats = dict(stats2)
+        stats["method"] = "cosine_with_tfidf"
+        stats["entropy"] = entropy
+        stats["tau"] = float(tau)
+        stats["tfidf_weight_mean"] = float(neuron_weights.mean().item())
+        stats["tfidf_weight_max"] = float(neuron_weights.max().item())
+        stats["initial_num_selected"] = int(first_mask_by_concept[concept].sum().item())
+
+        out[concept] = {
+            # Save weighted cosine as the primary "raw_cosine_similarities" output
+            # for downstream compatibility with existing plotting scripts.
+            "raw_cosine_similarities": weighted_cos,
+            # Keep original pre-TF-IDF cosine values for debugging/ablation.
+            "original_raw_cosine_similarities": raw_cos,
+            "weighted_cosine_similarities": weighted_cos,
+            "masked_cosine": masked_weighted,
+            "sae_text_activations": sae_by_concept[concept],
+            "stats": stats,
+        }
+
+    return out
+
+def build_concept_extractor(extractor_cfg: DictConfig):
+    """Resolve extractor callable from config and apply partial kwargs."""
+    if extractor_cfg is None:
+        raise ValueError("Missing concepts_experiment extractor config.")
+
+    target = extractor_cfg.get("target")
+    if target is None:
+        # Backward compatibility with older config format.
+        target = extractor_cfg.get("_extract_concept_distribution")
+    if not target:
+        raise ValueError(
+            "Extractor target missing. Set concepts_experiment.extractor.target "
+            "(or legacy methods._extract_concept_distribution)."
+        )
+
+    module_name, fn_name = str(target).rsplit(".", 1)
+    fn = getattr(importlib.import_module(module_name), fn_name)
+    if not callable(fn):
+        raise TypeError(f"Resolved extractor is not callable: {target}")
+
+    params_cfg = extractor_cfg.get("params", {})
+    params = dict(params_cfg) if params_cfg is not None else {}
+    return partial(fn, **params)
+
+
+
+def _assert_checkpoint_weights_loaded(model) -> None:
+    report = getattr(model, "_checkpoint_load_report", None)
+    if report is None:
+        raise RuntimeError(
+            "Checkpoint load report missing on model. "
+            "Ensure model is created via instantiate_model_and_load()."
+        )
+    if int(report.get("matched_keys", 0)) <= 0:
+        raise RuntimeError("Checkpoint weights do not match model parameters (0 matched keys).")
+    log.info(
+        "Checkpoint load verification passed: "
+        f"matched={report['matched_keys']}/{report['model_keys']}, "
+        f"missing={report['missing_keys']}, unexpected={report['unexpected_keys']}."
+    )
+
+
+def concept_from_bridge_score(
+    *,
+    model,
+    concept: str,
+    concept_index: int,
+    device: torch.device,
+    bridge_checkpoint_path: str = None,
+    bridge_checkpoint_key: str = "bridges",
+    dataloader=None,
+) -> torch.Tensor:
+    del concept_index  # kept for a stable extractor callable signature
+
+    def _extract_bridge_from_checkpoint(path: str, key: str):
+        checkpoint = load_state_dict_any(path, map_location="cpu")
+        state_dict = checkpoint.get("state_dict", checkpoint) if isinstance(checkpoint, dict) else checkpoint
+        if not isinstance(state_dict, dict):
+            return None
+
+        bridge_obj = state_dict.get(key)
+        if bridge_obj is None:
+            return None
+
+        if isinstance(bridge_obj, torch.Tensor):
+            return bridge_obj
+
+        if isinstance(bridge_obj, dict):
+            bridge_tensors = []
+            for _, val in sorted(bridge_obj.items(), key=lambda kv: str(kv[0])):
+                # New structure: bridges[dataset_name] = {"align": ..., "coact": ..., "B": ...}
+                if isinstance(val, dict):
+                    b_val = val.get("B")
+                    if isinstance(b_val, torch.Tensor):
+                        bridge_tensors.append(b_val)
+                # Fallback: dict values are directly bridge tensors.
+                elif isinstance(val, torch.Tensor):
+                    bridge_tensors.append(val)
+            if bridge_tensors:
+                if len(bridge_tensors) == 1:
+                    return bridge_tensors[0]
+                return torch.stack(bridge_tensors, dim=0).mean(dim=0)
+        return None
+
+    @torch.no_grad()
+    def _compute_bridge_from_dataloader(dataloader_obj) -> torch.Tensor:
+        za_chunks = []
+        zt_chunks = []
+        for batch in dataloader_obj:
+            audio = batch.get("audio") if isinstance(batch, dict) else None
+            text = batch.get("prompt") if isinstance(batch, dict) else None
+            if audio is None or text is None:
+                continue
+
+            audio = audio.to(device)
+            encoded_audio = (
+                model.audio_encoder(audio) if not getattr(model, "preextracted_features", False) else audio
+            )
+            encoded_text = model.text_encoder(text)
+            _, z_audio, _, _ = model(encoded_audio)
+            _, z_text, _, _ = model(encoded_text)
+            n = min(z_audio.shape[0], z_text.shape[0])
+            if n <= 0:
+                continue
+            za_chunks.append(z_audio[:n].detach())
+            zt_chunks.append(z_text[:n].detach())
+
+        if not za_chunks or not zt_chunks:
+            raise ValueError(
+                "Could not compute bridge from dataloader: no batches contained both audio and prompt."
+            )
+
+        za = torch.cat(za_chunks, dim=0)
+        zt = torch.cat(zt_chunks, dim=0)
+        w_dec = get_dictionary_from_lightningsae(model).to(device)
+        return BridgeScoreCallback._compute(za, zt, w_dec)["B"]
+
+    bridge_B = None
+    if bridge_checkpoint_path:
+        bridge_B = _extract_bridge_from_checkpoint(bridge_checkpoint_path, bridge_checkpoint_key)
+
+    if bridge_B is None:
+        if dataloader is None:
+            raise ValueError(
+                "bridges not found in checkpoint (or checkpoint not provided). "
+                "Provide dataloader to compute bridge scores on the fly."
+            )
+        bridge_B = _compute_bridge_from_dataloader(dataloader)
+
+    bridge_B = bridge_B.to(device)
+    text_embedding = model.text_encoder([concept])
+    activations = model.sae_encoder(text_embedding)
+    bridged_activations = bridge_B @ activations
+    return bridged_activations.detach().cpu()
+
+
+def run_extract_experiment(
+    *,
+    cfg: DictConfig,
+    model,
+    device: torch.device,
+    output_dir: str,
+    source_run_id: str,
+    ckpt_path: str,
+) -> Dict[str, Any]:
+
+
+    extractor_cfg = cfg.concepts_experiment.get("extractor")
+    if extractor_cfg is None:
+        extractor_cfg = cfg.concepts_experiment.get("methods", {})
+
+
+    extractor = build_concept_extractor(extractor_cfg)
+    extractor_target = str(extractor_cfg.get("target", ""))
+    concepts = read_vocab_lines(str(cfg.concepts.vocab_path))
+    _assert_checkpoint_weights_loaded(model)
+
+    out_ = {}
+
+    if extractor_target.endswith("cosine_with_tfidf"):
+        # Batch method needs all concepts to compute global neuron TF-IDF weights.
+        tfidf_results = extractor(
+            model=model,
+            concepts=concepts,
+            device=device,
+        )
+        for concept in concepts:
+            results = tfidf_results[concept]
+            masked_cosine = results["masked_cosine"].detach().cpu()
+            raw_cosine = results["raw_cosine_similarities"].detach().cpu()
+            weighted_cosine = results.get("weighted_cosine_similarities")
+            if isinstance(weighted_cosine, torch.Tensor):
+                weighted_cosine = weighted_cosine.detach().cpu()
+            sae_acts = results["sae_text_activations"]
+            if isinstance(sae_acts, torch.Tensor):
+                sae_acts = sae_acts.detach().cpu()
+            out_[concept] = {
+                "masked_cosine": masked_cosine,
+                "stats": results["stats"],
+                "raw_cosine_similarities": raw_cosine,
+                "weighted_cosine_similarities": weighted_cosine,
+                "sae_text_activations": sae_acts,
+            }
+    else:
+        for idx, concept in enumerate(concepts):
+            results = extractor(
+                model=model,
+                concept=concept,
+                concept_index=idx,
+                device=device,
+            )
+            masked_cosine = results["masked_cosine"].detach().cpu()
+            raw_cosine = results["raw_cosine_similarities"].detach().cpu()
+            sae_acts = results["sae_text_activations"]
+            if isinstance(sae_acts, torch.Tensor):
+                sae_acts = sae_acts.detach().cpu()
+            out_[concept] = {
+                "masked_cosine": masked_cosine,
+                "stats": results["stats"],
+                "raw_cosine_similarities": raw_cosine,
+                "sae_text_activations": sae_acts,
+            }
+
+    # Print a compact activation summary for quick inspection.
+    # "Active" means non-zero entries in the extracted activation tensor.
+    width = 40
+    topk_width = 20
+    summary_rows = []
+    total_active = 0
+    total_values = 0
+    global_max = None
+    for concept, values in out_.items():
+        tensor = values['masked_cosine'].detach().cpu()
+        if tensor.numel() == 0:
+            active = 0
+            size = 0
+            frac = 0.0
+            max_val = 0.0
+            topk_pairs = []
+        else:
+            active = int((tensor != 0).sum().item())
+            size = int(tensor.numel())
+            frac = active / max(size, 1)
+            max_val = float(tensor.max().item())
+            flat = tensor.flatten()
+            k = min(5, flat.numel())
+            top_vals = torch.topk(flat, k=k).values
+            denom = max(abs(max_val), 1e-12)
+            topk_pairs = []
+            for v in top_vals:
+                v_float = float(v.item())
+                rel = max(0.0, min(1.0, v_float / denom))
+                filled_top = int(round(rel * topk_width))
+                top_bar = "#" * filled_top + "-" * (topk_width - filled_top)
+                topk_pairs.append((v_float, top_bar))
+        if global_max is None:
+            global_max = max_val
+        else:
+            global_max = max(global_max, max_val)
+        total_active += active
+        total_values += size
+        filled = int(round(frac * width))
+        bar = "#" * filled + "-" * (width - filled)
+        summary_rows.append((concept, active, size, frac, max_val, bar, topk_pairs))
+
+    log.info("Activation summary per concept (active/total | active_frac | max):")
+    for concept, active, size, frac, max_val, bar, topk_pairs in summary_rows:
+        log.info(
+            f"- {concept:>24s} | {active:>7d}/{size:<7d} | {frac:6.2%} | "
+            f"max={max_val:>9.4f} | [{bar}]"
+        )
+        if topk_pairs:
+            topk_str = " | ".join([f"{val:9.4f} [{top_bar}]" for val, top_bar in topk_pairs])
+            log.info(f"  top5: {topk_str}")
+        else:
+            log.info("  top5: n/a")
+    global_frac = (total_active / total_values) if total_values > 0 else 0.0
+    if global_max is None:
+        global_max = 0.0
+    log.info(
+        "Activation summary total: "
+        f"{total_active}/{total_values} active ({global_frac:.2%}), global max={global_max:.4f}."
+    )
+    if out_:
+        for concept, values in out_.items():
+            stats = values['stats']
+            k95 = int(stats.get("k95", 0))
+            entropy = float(stats.get("entropy", 0.0))
+            log.info(
+                f"{concept:24s} | {k95:4d} | {entropy:9.4f}"
+            )
+
+
+    # Resolve output directory: <output_dir>/<run_id>/<experiment_id>
+    output_dir, run_id_slug, experiment_id = _resolve_output_dir(cfg, output_dir, source_run_id)
+
+    # Save activations under concepts.pt
+    activations_path = os.path.join(output_dir, "concepts.pt")
+    torch.save(out_, activations_path)
+    log.info(f"Saved concept activation dictionary to {activations_path}")
+
+    # Save config
+    config_path = os.path.join(output_dir, "config.yaml")
+    with open(config_path, "w") as f:
+        OmegaConf.save(cfg, f)
+
+
+    log.info(f"Saved concept activation dictionary to {activations_path}")
+
+    out_dict = {
+        "experiment": "concept_isolation",
+        "concepts_experiment": str(cfg.concepts_experiment.name),
+        "config_source": source_run_id,
+        "run_id": run_id_slug,
+        "experiment_id": experiment_id,
+        "checkpoint_path": ckpt_path,
+        "num_concepts": len(concepts),
+        "concept_activations_path": activations_path,
+    }
+
+    # Save JSON summary output
+    json_output_path = os.path.join(output_dir, "output.json")
+    import json
+    with open(json_output_path, "w") as f:
+        json.dump(out_dict, f)
+
+    return out_dict

steerable_retrieval/experiments/concepts/get_concepts.py

@@ -0,0 +1,497 @@
+#!/usr/bin/env python3
+"""
+Build a concept library for MuQ (MuQ-MuLan) from caption n-grams + groundedness on precomputed audio embeddings.
+
+Assumptions:
+- You have a local folder of ~5k audio embedding .npy files (each is [D] or [1,D] or [T,D] -> we pool to [D]).
+- You have a CSV on S3 with a column 'caption' containing text captions.
+- You have access to MuQMuLan via `from muq_mulan import MuQMuLan` (adjust import to your project).
+
+What it does:
+1) Load captions from s3://... CSV.
+2) Extract 1–3-gram candidates (stopword filtered + min docfreq).
+3) Embed all candidate phrases with MuQ text tower.
+4) Compute groundedness tail score against your audio embedding set.
+5) Select a diverse set via greedy MMR (optional).
+6) Save CSV/JSONL with concepts + stats.
+
+Example:
+python build_concept_library.py \
+  --audio_dir /path/to/muq_audio_npys \
+  --captions_s3 s3://my-bucket/path/captions.csv \
+  --out_csv concepts.csv \
+  --max_candidates 5000 \
+  --final_k 250 \
+  --device cuda
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import glob
+import json
+import math
+import os
+import re
+import warnings
+from dataclasses import dataclass
+from typing import Dict, Iterable, List, Optional, Tuple
+
+import numpy as np
+import torch
+
+# -----------------------------
+# Text encoder (given)
+# -----------------------------
+# IMPORTANT: adjust imports to your codebase.
+# from openmuq import MuQMuLan  # <-- likely incorrect; replace with your actual import
+from steerable_retrieval.models.encoders.muq import MuQTextEncoder
+
+# -----------------------------
+# S3 CSV loader
+# -----------------------------
+def read_captions_from_s3_csv(s3_uri: str, caption_col: str = "caption", max_rows: Optional[int] = None) -> List[str]:
+    """
+    Read captions from an S3 CSV using boto3 streaming.
+    Requires AWS credentials in environment / config.
+
+    s3_uri: s3://bucket/key.csv
+    """
+    import boto3
+    from botocore.config import Config
+
+    m = re.match(r"^s3://([^/]+)/(.+)$", s3_uri)
+    if not m:
+        raise ValueError(f"Invalid s3 uri: {s3_uri}")
+    bucket, key = m.group(1), m.group(2)
+
+    s3 = boto3.client("s3", config=Config(signature_version="s3v4"))
+    obj = s3.get_object(Bucket=bucket, Key=key)
+    body = obj["Body"]
+
+    # Stream decode bytes -> text lines
+    lines = (line.decode("utf-8", errors="replace") for line in body.iter_lines())
+
+    reader = csv.DictReader(lines)
+    caps: List[str] = []
+    for i, row in enumerate(reader):
+        if caption_col not in row:
+            raise KeyError(f"CSV missing column '{caption_col}'. Columns: {list(row.keys())}")
+        cap = (row[caption_col] or "").strip()
+        if cap:
+            caps.append(cap)
+        if max_rows is not None and (i + 1) >= max_rows:
+            break
+    return caps
+
+
+# -----------------------------
+# Caption n-gram mining
+# -----------------------------
+_DEFAULT_STOPWORDS = set(
+    """
+a an the and or but if while with without within into onto of for in on at by from as
+is are was were be been being have has had do does did can could may might will would
+this that these those it its it's i you we they them our your their
+music song track sound sounds audio instrumental vocals voice listen listening
+featuring feat ft version remix edit mix original
+genre genres follows follow under includes include including characterized
+scenario scenarios scene scenes depicting suitable settings overall
+creates creating evokes around approximately set one main piece
+tempo bpm key signature time
+0 1 2 3 4 5 6 7 8 9
+""".split()
+)
+
+_TOKEN_RE = re.compile(r"[a-z0-9]+(?:'[a-z0-9]+)?", re.IGNORECASE)
+
+
+def _get_merged_stopwords() -> set:
+    """Return built-in stopwords merged with NLTK English stopwords when available."""
+    stopwords = set(_DEFAULT_STOPWORDS)
+    try:
+        from nltk.corpus import stopwords as nltk_stopwords  # type: ignore
+        try:
+            stopwords.update(nltk_stopwords.words("english"))
+        except LookupError:
+            import nltk  # type: ignore
+
+            nltk.download("stopwords", quiet=True)
+            stopwords.update(nltk_stopwords.words("english"))
+    except Exception as exc:
+        warnings.warn(f"NLTK stopwords unavailable, using built-in stopwords only: {exc}")
+    return stopwords
+
+
+def normalize_text(s: str) -> str:
+    s = s.lower()
+    s = re.sub(r"\s+", " ", s).strip()
+    return s
+
+
+def tokenize(s: str) -> List[str]:
+    return _TOKEN_RE.findall(s.lower())
+
+
+def extract_ngrams(tokens: List[str], n_min: int = 1, n_max: int = 3) -> List[str]:
+    out: List[str] = []
+    L = len(tokens)
+    for n in range(n_min, n_max + 1):
+        for i in range(0, L - n + 1):
+            out.append(" ".join(tokens[i : i + n]))
+    return out
+
+
+def mine_ngram_candidates(
+    captions: List[str],
+    n_min: int = 1,
+    n_max: int = 3,
+    stopwords: Optional[set] = None,
+    min_df: int = 5,
+    max_df_frac: float = 0.30,
+    max_candidates: int = 5000,
+) -> List[Tuple[str, int]]:
+    """
+    Returns list of (ngram, df) sorted by df desc, then length desc.
+    Uses document frequency: presence in a caption at least once.
+    """
+    if stopwords is None:
+        stopwords = _get_merged_stopwords()
+
+    df: Dict[str, int] = {}
+    N = len(captions)
+
+    for cap in captions:
+        toks = [t for t in tokenize(cap) if t not in stopwords]
+        if not toks:
+            continue
+        grams = set(extract_ngrams(toks, n_min, n_max))
+        for g in grams:
+            df[g] = df.get(g, 0) + 1
+
+    # filter by df
+    max_df = int(max_df_frac * N)
+    items = [(g, c) for g, c in df.items() if c >= min_df and c <= max_df]
+
+    # sort by df, then prefer longer phrases
+    items.sort(key=lambda x: (x[1], len(x[0].split())), reverse=True)
+
+    if len(items) > max_candidates:
+        items = items[:max_candidates]
+    return items
+
+
+# -----------------------------
+# Audio embedding loading
+# -----------------------------
+def load_audio_embeddings(audio_dir: str, max_files: Optional[int] = None) -> np.ndarray:
+    """
+    Loads .npy files, returns array [N, D].
+    Pools if embedding has shape [T, D] by mean over T.
+    """
+    paths = sorted(glob.glob(os.path.join(audio_dir, "*.npy")))
+    if not paths:
+        raise FileNotFoundError(f"No .npy files found in {audio_dir}")
+    if max_files is not None:
+        paths = paths[:max_files]
+
+    embs: List[np.ndarray] = []
+    D: Optional[int] = None
+    for p in paths:
+        x = np.load(p)
+        x = np.asarray(x)
+        if x.ndim == 2:
+            # [1, D] or [T, D]
+            if x.shape[0] == 1:
+                x = x[0]
+            else:
+                x = x.mean(axis=0)
+        elif x.ndim == 1:
+            pass
+        else:
+            raise ValueError(f"Unexpected embedding shape {x.shape} in {p}")
+
+        if D is None:
+            D = int(x.shape[0])
+        elif int(x.shape[0]) != D:
+            raise ValueError(f"Dimension mismatch: got {x.shape[0]} vs expected {D} in {p}")
+
+        embs.append(x.astype(np.float32))
+
+    A = np.stack(embs, axis=0)  # [N, D]
+    return A
+
+
+def l2_normalize_np(x: np.ndarray, eps: float = 1e-8) -> np.ndarray:
+    n = np.linalg.norm(x, axis=-1, keepdims=True)
+    return x / (n + eps)
+
+
+# -----------------------------
+# Groundedness scoring
+# -----------------------------
+@dataclass
+class GroundScore:
+    concept: str
+    df: int
+    mean: float
+    topq_mean: float
+    tail_gap: float
+    gini: float
+    pos_mass: float  # fraction of sims > 0
+
+
+def gini_coefficient(x: np.ndarray) -> float:
+    """
+    Gini for nonnegative values; for similarities, we shift to nonnegative.
+    """
+    x = np.asarray(x, dtype=np.float64)
+    if x.size == 0:
+        return 0.0
+    if np.allclose(x, 0):
+        return 0.0
+    x = np.sort(x)
+    n = x.size
+    cum = np.cumsum(x)
+    # Gini = 1 - 2 * sum_i (cum_i) / (n * cum_n) + 1/n
+    return float(1.0 - (2.0 * np.sum(cum) / (n * cum[-1] + 1e-12)) + (1.0 / n))
+
+
+@torch.no_grad()
+def embed_text_batch(
+    enc: MuQTextEncoder, texts: List[str], batch_size: int = 64, device: str = "cuda"
+) -> np.ndarray:
+    feats: List[np.ndarray] = []
+    for i in range(0, len(texts), batch_size):
+        batch = texts[i : i + batch_size]
+        t = enc(batch).to(device)
+        t = t.float()
+        t = torch.nn.functional.normalize(t, dim=-1)
+        feats.append(t.detach().cpu().numpy())
+    return np.concatenate(feats, axis=0)
+
+
+def compute_groundedness(
+    text_embs: np.ndarray,  # [C, D] L2-normalized
+    audio_embs: np.ndarray,  # [N, D] L2-normalized
+    concepts: List[str],
+    dfs: List[int],
+    top_q: float = 0.01,
+) -> List[GroundScore]:
+    """
+    For each concept embedding, compute similarity distribution vs audio embeddings and derive scores.
+    """
+    A = audio_embs
+    C = text_embs
+    assert C.shape[0] == len(concepts) == len(dfs)
+
+    scores: List[GroundScore] = []
+
+    # Compute in chunks to avoid huge memory if C is large.
+    # sims = A @ C.T -> [N, C]
+    # With N=5000, C up to 5000, this is 25M ~ 100MB float32; still OK, but chunk anyway.
+    N = A.shape[0]
+    qk = max(1, int(math.ceil(top_q * N)))
+
+    chunk = 512
+    for j0 in range(0, C.shape[0], chunk):
+        j1 = min(C.shape[0], j0 + chunk)
+        Cc = C[j0:j1]  # [c, D]
+        sims = A @ Cc.T  # [N, c]
+
+        for jj in range(j1 - j0):
+            s = sims[:, jj].astype(np.float64)
+            mean = float(s.mean())
+            # top-q mean
+            top_idx = np.argpartition(s, -qk)[-qk:]
+            topq_mean = float(s[top_idx].mean())
+            tail_gap = topq_mean - mean
+
+            # groundedness also reflected by concentration: compute gini on shifted sims
+            s_shift = s - s.min()
+            gini = gini_coefficient(s_shift)
+
+            pos_mass = float((s > 0).mean())
+
+            scores.append(
+                GroundScore(
+                    concept=concepts[j0 + jj],
+                    df=dfs[j0 + jj],
+                    mean=mean,
+                    topq_mean=topq_mean,
+                    tail_gap=tail_gap,
+                    gini=gini,
+                    pos_mass=pos_mass,
+                )
+            )
+    return scores
+
+
+# -----------------------------
+# Diversity selection (MMR)
+# -----------------------------
+def mmr_select(
+    concepts: List[str],
+    text_embs: np.ndarray,  # [C, D] L2-normalized
+    base_score: np.ndarray,  # [C]
+    k: int,
+    lambda_mmr: float = 0.7,
+) -> List[int]:
+    """
+    Greedy Maximal Marginal Relevance selection.
+    """
+    C = text_embs
+    selected: List[int] = []
+    remaining = set(range(len(concepts)))
+
+    # Precompute similarity matrix chunk-wise if needed; for C <= 5000 OK but keep simple:
+    # We'll compute max similarity to selected on the fly.
+    while len(selected) < k and remaining:
+        best_i = None
+        best_val = -1e18
+        for i in list(remaining):
+            if not selected:
+                val = float(base_score[i])
+            else:
+                # max similarity to selected
+                sims = C[selected] @ C[i]
+                max_sim = float(np.max(sims))
+                val = lambda_mmr * float(base_score[i]) - (1.0 - lambda_mmr) * max_sim
+            if val > best_val:
+                best_val = val
+                best_i = i
+        assert best_i is not None
+        selected.append(best_i)
+        remaining.remove(best_i)
+
+    return selected
+
+
+# -----------------------------
+# Main
+# -----------------------------
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--audio_dir", type=str, required=True)
+    ap.add_argument("--captions_s3", type=str, required=True)
+    ap.add_argument("--caption_col", type=str, default="caption")
+    ap.add_argument("--max_captions", type=int, default=10000)
+
+    ap.add_argument("--n_min", type=int, default=1)
+    ap.add_argument("--n_max", type=int, default=3)
+    ap.add_argument("--min_df", type=int, default=5)
+    ap.add_argument("--max_df_frac", type=float, default=0.30)
+    ap.add_argument("--max_candidates", type=int, default=5000)
+
+    ap.add_argument("--model_name", type=str, default="OpenMuQ/MuQ-MuLan-large")
+    ap.add_argument("--device", type=str, default="cpu")
+    ap.add_argument("--text_batch_size", type=int, default=16)
+
+    ap.add_argument("--top_q", type=float, default=0.01)
+    ap.add_argument("--final_k", type=int, default=20)
+    ap.add_argument("--use_mmr", action="store_true")
+    ap.add_argument("--lambda_mmr", type=float, default=0.7)
+
+    ap.add_argument("--out_csv", type=str, required=True)
+    ap.add_argument("--out_jsonl", type=str, default=None)
+
+    args = ap.parse_args()
+
+    print(f"[1/6] Loading audio embeddings from {args.audio_dir} ...")
+    audio = load_audio_embeddings(args.audio_dir)
+    audio = l2_normalize_np(audio)
+    print(f"Loaded audio embeddings: {audio.shape}")
+
+    print(f"[2/6] Reading captions from {args.captions_s3} ...")
+    captions = read_captions_from_s3_csv(args.captions_s3, caption_col=args.caption_col, max_rows=args.max_captions)
+    print(f"Loaded captions: {len(captions)}")
+
+    print("[3/6] Mining n-gram candidates ...")
+    cand = mine_ngram_candidates(
+        captions=captions,
+        n_min=args.n_min,
+        n_max=args.n_max,
+        min_df=args.min_df,
+        max_df_frac=args.max_df_frac,
+        max_candidates=args.max_candidates,
+    )
+    concepts = [g for g, _ in cand]
+    dfs = [df for _, df in cand]
+    print(f"Candidates after filtering: {len(concepts)}")
+
+    print("[4/6] Embedding candidate concepts with MuQ text encoder ...")
+    enc = MuQTextEncoder(model_name=args.model_name, device=args.device, freeze=True)
+    text_embs = embed_text_batch(enc, concepts, batch_size=args.text_batch_size, device=args.device)
+    text_embs = l2_normalize_np(text_embs)
+    print(f"Text embeddings: {text_embs.shape}")
+
+    print("[5/6] Computing groundedness scores ...")
+    scores = compute_groundedness(
+        text_embs=text_embs,
+        audio_embs=audio,
+        concepts=concepts,
+        dfs=dfs,
+        top_q=args.top_q,
+    )
+
+    # Base score for ranking: tail_gap (primary) * log(df+1) (secondary) + small gini
+    base = np.array([s.tail_gap for s in scores], dtype=np.float64)
+    dfw = np.log(np.array([s.df for s in scores], dtype=np.float64) + 1.0)
+    gini = np.array([s.gini for s in scores], dtype=np.float64)
+    base_score = base * dfw + 0.05 * gini
+
+    # Sort by base_score
+    order = np.argsort(-base_score)
+    if args.use_mmr:
+        # Run MMR on top pool to keep it fast
+        pool = order[: min(len(order), max(args.final_k * 10, 500))]
+        pool_concepts = [concepts[i] for i in pool]
+        pool_embs = text_embs[pool]
+        pool_base = base_score[pool]
+        sel_pool_idx = mmr_select(pool_concepts, pool_embs, pool_base, k=args.final_k, lambda_mmr=args.lambda_mmr)
+        selected = [pool[i] for i in sel_pool_idx]
+    else:
+        selected = order[: args.final_k].tolist()
+
+    print(f"[6/6] Writing outputs: top-{len(selected)} concepts -> {args.out_csv}")
+    os.makedirs(os.path.dirname(args.out_csv) or ".", exist_ok=True)
+    # also write the concepts to a text file
+    with open(args.out_csv.replace(".csv", ".txt"), "w", encoding="utf-8") as f:
+        for i in selected:
+            f.write(scores[i].concept + "\n")
+
+    with open(args.out_csv, "w", newline="", encoding="utf-8") as f:
+        w = csv.writer(f)
+        w.writerow(["concept", "df", "mean_sim", "topq_mean_sim", "tail_gap", "gini", "pos_mass", "base_score"])
+        for i in selected:
+            s = scores[i]
+            w.writerow([s.concept, s.df, s.mean, s.topq_mean, s.tail_gap, s.gini, s.pos_mass, float(base_score[i])])
+
+    if args.out_jsonl:
+        os.makedirs(os.path.dirname(args.out_jsonl) or ".", exist_ok=True)
+        with open(args.out_jsonl, "w", encoding="utf-8") as f:
+            for i in selected:
+                s = scores[i]
+                f.write(
+                    json.dumps(
+                        {
+                            "concept": s.concept,
+                            "df": s.df,
+                            "mean_sim": s.mean,
+                            "topq_mean_sim": s.topq_mean,
+                            "tail_gap": s.tail_gap,
+                            "gini": s.gini,
+                            "pos_mass": s.pos_mass,
+                            "base_score": float(base_score[i]),
+                        },
+                        ensure_ascii=False,
+                    )
+                    + "\n"
+                )
+
+    print("Done.")
+
+
+if __name__ == "__main__":
+    main()

steerable_retrieval/experiments/notebook_cache.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+import hashlib
+import json
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import pandas as pd
+from omegaconf import OmegaConf
+
+try:
+    import torch
+except Exception:  # pragma: no cover
+    torch = None
+
+
+def _normalize_for_hash(value: Any) -> Any:
+    if OmegaConf.is_config(value):
+        value = OmegaConf.to_container(value, resolve=True)
+
+    if isinstance(value, dict):
+        return {str(key): _normalize_for_hash(val) for key, val in sorted(value.items(), key=lambda item: str(item[0]))}
+    if isinstance(value, (list, tuple)):
+        return [_normalize_for_hash(item) for item in value]
+    if isinstance(value, set):
+        return sorted(_normalize_for_hash(item) for item in value)
+    if isinstance(value, Path):
+        return str(value)
+    if isinstance(value, np.ndarray):
+        return value.tolist()
+    if isinstance(value, np.generic):
+        return value.item()
+    if torch is not None and isinstance(value, torch.Tensor):
+        return value.detach().cpu().tolist()
+    if torch is not None and isinstance(value, torch.device):
+        return str(value)
+    if isinstance(value, (str, int, float, bool)) or value is None:
+        return value
+    return repr(value)
+
+
+def stable_hash(value: Any, length: int = 16) -> str:
+    payload = json.dumps(
+        _normalize_for_hash(value),
+        ensure_ascii=True,
+        separators=(",", ":"),
+        sort_keys=True,
+    )
+    return hashlib.sha256(payload.encode("utf-8")).hexdigest()[:length]
+
+
+def build_cache_path(
+    cache_root: str | Path,
+    eval_config: Any,
+    run_id: str,
+    model_config: Any,
+    suffix: str = "edit.pkl",
+) -> tuple[str, str, Path]:
+    eval_hash = stable_hash(eval_config)
+    model_hash = stable_hash({"run_id": run_id, "model_config": model_config})
+    cache_path = Path(cache_root) / f"{eval_hash}_{model_hash}_{suffix}"
+    return eval_hash, model_hash, cache_path
+
+
+def load_cached_frame(cache_path: str | Path) -> pd.DataFrame:
+    return pd.read_pickle(Path(cache_path))
+
+
+def save_cached_frame(df: pd.DataFrame, cache_path: str | Path) -> Path:
+    cache_path = Path(cache_path)
+    cache_path.parent.mkdir(parents=True, exist_ok=True)
+    df.to_pickle(cache_path)
+    return cache_path
+
+
+def build_cache_metadata_path(cache_path: str | Path, metadata_filename: str = "hash.json") -> Path:
+    cache_path = Path(cache_path)
+    cache_stem = cache_path.name
+    if cache_path.suffix:
+        cache_stem = cache_path.name[: -len(cache_path.suffix)]
+    return cache_path.parent / f"{cache_stem}.{metadata_filename}"
+
+
+def save_cache_metadata(metadata: Any, cache_path: str | Path, metadata_filename: str = "hash.json") -> Path:
+    metadata_path = build_cache_metadata_path(cache_path, metadata_filename=metadata_filename)
+    metadata_path.parent.mkdir(parents=True, exist_ok=True)
+    with metadata_path.open("w", encoding="utf-8") as f:
+        json.dump(_normalize_for_hash(metadata), f, indent=2, sort_keys=True)
+    return metadata_path

steerable_retrieval/experiments/run_concept_isolation.py

@@ -0,0 +1,71 @@
+import json
+import os
+from typing import Any, Dict, Optional
+
+import lightning as L
+import rootutils
+import torch
+from dora.hydra import hydra_main
+from omegaconf import DictConfig
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+from steerable_retrieval.experiments.common import (  # noqa: E402
+    ensure_dir,
+    instantiate_model_and_load,
+    resolve_model_from_source,
+    save_config,
+)
+from steerable_retrieval.experiments.concepts.extract import run_extract_experiment
+from steerable_retrieval.utils import RankedLogger, extras, register_resolvers  # noqa: E402
+
+log = RankedLogger(__name__, rank_zero_only=True)
+register_resolvers()
+
+
+def run_concept_isolation(cfg: DictConfig) -> Dict[str, Any]:
+    if cfg.get("seed") is not None:
+        L.seed_everything(int(cfg.seed), workers=True)
+
+    source_run_id = resolve_model_from_source(cfg)
+    output_dir = ensure_dir(str(cfg.output_dir))
+    save_config(cfg, output_dir)
+
+    device = torch.device(cfg.get("device", "cuda:0") if torch.cuda.is_available() else "cpu")
+    log.info(f"Using device: {device}")
+
+    model, ckpt_path = instantiate_model_and_load(cfg, device=device)
+    experiment_name = str(cfg.concepts_experiment.get("name", "extract"))
+
+    if experiment_name == "extract":
+        out_dict = run_extract_experiment(
+            cfg=cfg,
+            model=model,
+            device=device,
+            output_dir=output_dir,
+            source_run_id=source_run_id,
+            ckpt_path=ckpt_path,
+        )
+    else:
+        raise ValueError(
+            f"Unsupported concepts_experiment.name={experiment_name}. "
+            "Expected one of: extract, score_edit_add, score_edit_suppress."
+        )
+
+    # Save metadata for downstream launch tooling.
+    with open(os.path.join(output_dir, "output.json"), "w") as f:
+        json.dump(out_dict, f)
+
+    return out_dict
+
+
+@hydra_main(version_base="1.3", config_path="../../configs", config_name="experiment/concept_isolation.yaml")
+def main(cfg: DictConfig) -> Optional[Dict[str, Any]]:
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
+    extras(cfg)
+    return run_concept_isolation(cfg)
+
+
+if __name__ == "__main__":
+    main()
+

steerable_retrieval/experiments/run_stability.py

@@ -0,0 +1,152 @@
+import copy
+import math
+import os
+from typing import Any, Dict, List, Optional
+
+import lightning as L
+import rootutils
+import torch
+import torch.nn.functional as F
+from dora.hydra import hydra_main
+from omegaconf import DictConfig
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+from steerable_retrieval.callbacks.energy import get_dictionary_from_lightningsae, load_state_dict_any  # noqa: E402
+from steerable_retrieval.experiments.common import (  # noqa: E402
+    ensure_dir,
+    finish_wandb,
+    instantiate_model_and_load,
+    maybe_start_wandb,
+    resolve_model_from_source,
+    save_config,
+    save_json,
+)
+from steerable_retrieval.utils import RankedLogger, extras, register_resolvers  # noqa: E402
+
+log = RankedLogger(__name__, rank_zero_only=True)
+register_resolvers()
+
+
+def _stability_hungarian(a: torch.Tensor, b: torch.Tensor) -> float:
+    from scipy.optimize import linear_sum_assignment
+
+    sim = (a @ b.t()).detach().cpu().numpy()
+    row_idx, col_idx = linear_sum_assignment(-sim)
+    return float(sim[row_idx, col_idx].mean())
+
+
+def _k_values(n_concepts: int, n_curve_points: int) -> List[int]:
+    if n_concepts <= 1:
+        return [1]
+    out = [1]
+    for i in range(max(2, n_curve_points)):
+        frac = i / max(n_curve_points - 1, 1)
+        k = int(round(math.exp(math.log(1) * (1 - frac) + math.log(n_concepts) * frac)))
+        out.append(max(1, min(n_concepts, k)))
+    out.append(n_concepts)
+    return sorted(set(out))
+
+
+def run_stability(cfg: DictConfig) -> Dict[str, Any]:
+    if cfg.get("seed") is not None:
+        L.seed_everything(int(cfg.seed), workers=True)
+
+    source_run_id = resolve_model_from_source(cfg)
+    output_dir = ensure_dir(str(cfg.output_dir))
+    save_config(cfg, output_dir)
+
+    device = torch.device(cfg.get("device", "cuda:0") if torch.cuda.is_available() else "cpu")
+    model, ckpt_path = instantiate_model_and_load(cfg, device=device)
+    d0 = F.normalize(get_dictionary_from_lightningsae(model).to(device), dim=-1)
+    n_concepts = int(d0.shape[0])
+
+    checkpoints = list(cfg.stability.checkpoint_paths)
+    if not checkpoints:
+        raise ValueError("stability.checkpoint_paths cannot be empty.")
+
+    run = maybe_start_wandb(
+        cfg=cfg,
+        job_type="stability_processing",
+        name=str(cfg.get("run_name", "stability_processing")),
+        output_dir=output_dir,
+    )
+
+    sorted_indices = torch.arange(n_concepts, device=device)
+    k_values = _k_values(n_concepts=n_concepts, n_curve_points=int(cfg.stability.n_curve_points))
+    topk = cfg.stability.get("topk_by_energy")
+
+    per_checkpoint = []
+    all_stabilities = []
+    all_topk = []
+    curves = []
+
+    for path in checkpoints:
+        comparison = copy.deepcopy(model).to("cpu")
+        comparison.load_state_dict(load_state_dict_any(path, map_location="cpu"), strict=False)
+        comparison = comparison.to(device).eval()
+        dk = F.normalize(get_dictionary_from_lightningsae(comparison).to(device), dim=-1)
+
+        stab_full = _stability_hungarian(d0, dk)
+        all_stabilities.append(stab_full)
+
+        stab_topk = None
+        if topk is not None:
+            k = min(int(topk), n_concepts)
+            top_idx = sorted_indices[:k]
+            stab_topk = _stability_hungarian(d0[top_idx], dk[top_idx])
+            all_topk.append(stab_topk)
+
+        curve_vals = []
+        for k in k_values:
+            top_idx = sorted_indices[:k]
+            curve_vals.append(_stability_hungarian(d0[top_idx], dk[top_idx]))
+        curves.append(curve_vals)
+
+        per_checkpoint.append(
+            {
+                "checkpoint_path": str(path),
+                "stability_full": stab_full,
+                "stability_topk": stab_topk,
+            }
+        )
+
+    mean_curve = [sum(vals) / len(vals) for vals in zip(*curves)]
+    summary = {
+        "experiment": "stability_processing",
+        "config_source_wandb_run_id": source_run_id,
+        "source_checkpoint": ckpt_path,
+        "n_concepts": n_concepts,
+        "per_checkpoint": per_checkpoint,
+        "mean_stability_full": float(sum(all_stabilities) / len(all_stabilities)),
+        "mean_stability_topk": float(sum(all_topk) / len(all_topk)) if all_topk else None,
+        "k_values": k_values,
+        "mean_stability_curve": mean_curve,
+    }
+    save_json(os.path.join(output_dir, "summary.json"), summary)
+
+    if run is not None:
+        import wandb
+
+        table = wandb.Table(columns=["checkpoint_path", "stability_full", "stability_topk"])
+        for item in per_checkpoint:
+            table.add_data(item["checkpoint_path"], item["stability_full"], item["stability_topk"])
+        wandb.log({"stability/per_checkpoint": table})
+        wandb.log({"stability/mean_full": summary["mean_stability_full"]})
+        if summary["mean_stability_topk"] is not None:
+            wandb.log({"stability/mean_topk": summary["mean_stability_topk"]})
+
+    finish_wandb(run)
+    return summary
+
+
+@hydra_main(version_base="1.3", config_path="../../configs/experiment", config_name="stability.yaml")
+def main(cfg: DictConfig) -> Optional[Dict[str, Any]]:
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
+    extras(cfg)
+    return run_stability(cfg)
+
+
+if __name__ == "__main__":
+    main()
+

steerable_retrieval/extract_dataset.py

@@ -0,0 +1,295 @@
+import os
+from typing import Any, Dict, Optional, Set
+from datetime import datetime
+
+import lightning as L
+import rootutils
+import torch
+import hydra
+from omegaconf import DictConfig, OmegaConf
+
+
+from dora import get_xp, hydra_main
+import hydra
+
+import lightning as L
+import rootutils
+import torch
+from omegaconf import DictConfig
+
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+# ------------------------------------------------------------------------------------ #
+# the setup_root above is equivalent to:
+# - adding project root dir to PYTHONPATH
+#       (so you don't need to force user to install project as a package)
+#       (necessary before importing any local modules e.g. `from gdr import utils`)
+# - setting up PROJECT_ROOT environment variable
+#       (which is used as a base for paths in "configs/paths/default.yaml")
+#       (this way all filepaths are the same no matter where you run the code)
+# - loading environment variables from ".env" in root dir
+#
+# you can remove it if you:
+# 1. either install project as a package or move entry files to project root dir
+# 2. set `root_dir` to "." in "configs/paths/default.yaml"
+#
+# more info: https://github.com/ashleve/rootutils
+# ------------------------------------------------------------------------------------ #
+
+from steerable_retrieval.utils import (
+    RankedLogger,
+    extras,
+    register_resolvers,
+)
+
+log = RankedLogger(__name__, rank_zero_only=True)
+register_resolvers()
+
+
+def list_existing_paths(path: str) -> Set[str]:
+    """List existing feature files from either a local directory or S3 path.
+    
+    Args:
+        path: Either a local directory path or S3 path (s3://bucket/path) where features are stored
+        
+    Returns:
+        Set of relative file paths (done IDs) extracted from existing files
+    """
+    done_ids = set()
+    
+    if not path:
+        log.info("No path provided, skipping path listing")
+        return done_ids
+    
+    # Check if it's an S3 path
+    if 's3://' in path:
+        # S3 path - list from S3
+        try:
+            import boto3
+            client = boto3.client('s3')
+            
+            # Extract bucket and prefix from S3 path
+            s3_path = path.replace("s3://", "")
+            parts = s3_path.split("/", 1)
+            bucket = parts[0]
+            prefix = parts[1] if len(parts) > 1 else ""
+            if prefix and not prefix.endswith('/'):
+                prefix = prefix + '/'
+            
+            log.info(f"Listing existing paths from S3: s3://{bucket}/{prefix}")
+            
+            # Paginate through all objects in the target location
+            paginator = client.get_paginator('list_objects_v2')
+            page_iterator = paginator.paginate(Bucket=bucket, Prefix=prefix)
+            
+            path_count = 0
+            for page in page_iterator:
+                if 'Contents' in page:
+                    for obj in page['Contents']:
+                        object_key = obj['Key']
+                        # Extract relative path by removing the prefix
+                        if object_key.startswith(prefix):
+                            relative_path = object_key[len(prefix):]
+                            # Skip empty paths and directories (ending with /)
+                            if relative_path and not relative_path.endswith('/'):
+                                # Add the path as-is
+                                done_ids.add(relative_path.replace('.npy','').replace('.wav','').replace('.mp3',''))
+                                # remove extension
+                                path_count += 1
+            
+            log.info(f"Found {path_count} existing paths in S3, {len(done_ids)} unique done IDs")
+            
+        except Exception as e:
+            log.warning(f"Failed to list existing paths from S3: {e}")
+            log.warning("Continuing without skipping already processed items")
+    else:
+        # Local directory - list from filesystem
+        try:
+            if not os.path.exists(path):
+                log.info(f"Local path does not exist: {path}, skipping path listing")
+                return done_ids
+            
+            log.info(f"Listing existing paths from local directory: {path}")
+            
+            # Walk through the directory and collect all .npy files
+            path_count = 0
+            for root, dirs, files in os.walk(path):
+                for file in files:
+                    if file.endswith('.npy'):
+                        # Get relative path from the base path
+                        full_path = os.path.join(root, file)
+                        relative_path = os.path.relpath(full_path, path)
+                        # Normalize path separators (use forward slashes)
+                        relative_path = relative_path.replace(os.sep, '/')
+                        
+                        # Add the path as-is
+                        done_ids.add(relative_path.replace('.npy','').replace('.wav','').replace('.mp3',''))
+                        path_count += 1
+            
+            log.info(f"Found {path_count} existing paths locally, {len(done_ids)} unique done IDs")
+            
+        except Exception as e:
+            log.warning(f"Failed to list existing paths from local directory: {e}")
+            log.warning("Continuing without skipping already processed items")
+    
+    return done_ids
+
+
+def extract_features(cfg: DictConfig) -> Dict[str, Any]:
+    """Extracts features from audio datasets using pre-trained encoders.
+
+    :param cfg: A DictConfig configuration composed by Hydra.
+    :return: A dict with extraction metadata.
+    """
+    # set seed for random number generators in pytorch, numpy and python.random
+    if cfg.get("seed"):
+        L.seed_everything(cfg.seed, workers=True)
+
+    log.info(f"Instantiating datamodule <{cfg.data._target_}>")
+    datamodule = hydra.utils.instantiate(cfg.data)
+
+    log.info(f"Instantiating model <{cfg.model._target_}>")
+    model = hydra.utils.instantiate(cfg.model)
+
+    # Setup datamodule
+    datamodule.setup(None)
+    
+    # Move encoder_pair to device
+    device = cfg.get("device", "cuda:0") if torch.cuda.is_available() else "cpu"
+    model.to(device)
+    log.info(f"Using device: {device}")
+
+
+    # Get extract parameters from config
+    save_dir = cfg.get("save_dir")
+    root_path = cfg.get("root_path")
+    extract_method = cfg.get("extract_method", "get_audio_embedding_from_data")
+    out_key = cfg.get("out_key", None)
+    hop = cfg.get("hop", 48000)
+    limit_n = cfg.get("limit_n")
+    save = cfg.get("save", False)
+    sagemaker_dir = cfg.get("sagemaker_dir", None)
+
+    # Handle save_dir
+    if save_dir is None:
+        if datamodule.train_dataset is not None and len(datamodule.train_dataset.annotations) > 0:
+            save_dir = os.path.dirname(datamodule.train_dataset.annotations[0]['file_path'])
+        else:
+            raise ValueError("save_dir must be specified in config or available from dataset")
+
+    # Check for sagemaker config to get S3 output destination
+    # In sagemaker, files are saved locally to save_dir, then uploaded to S3 destination
+    s3_destination = sagemaker_dir
+    log.info(f"S3 destination: {s3_destination}")
+    # List existing paths on restart to build done_ids set
+    # Check both save_dir and s3_destination (if using SageMaker) to avoid re-extracting
+    done_ids = set()
+    if save:
+        # Check save_dir (local or S3)
+        done_ids_save_dir = list_existing_paths(save_dir)
+        if done_ids_save_dir:
+            log.info(f"Found {len(done_ids_save_dir)} existing files in save_dir: {save_dir}")
+        done_ids.update(done_ids_save_dir)
+        
+        # Also check S3 destination if using SageMaker
+        if s3_destination:
+            done_ids_s3 = list_existing_paths(s3_destination)
+            if done_ids_s3:
+                log.info(f"Found {len(done_ids_s3)} existing files in S3 destination: {s3_destination}")
+            done_ids.update(done_ids_s3)
+        
+        if done_ids:
+            log.info(f"Total: Skipping {len(done_ids)} already processed items")
+        else:
+            log.info("No existing paths found in save_dir or S3 destination")
+
+
+    log.info(f"Extracting features with {extract_method} method")
+    log.info(f"Saving to: {save_dir}")
+    log.info(f"out_key: {out_key}, hop: {hop}, limit_n: {limit_n}, save: {save}")
+
+    # Save config to save_dir
+    if 's3://' in save_dir:
+        # S3 path - upload config
+        try:
+            import s3fs
+            fs = s3fs.S3FileSystem()
+            config_path = f"{save_dir}/config.yaml"
+            
+            with fs.open(config_path, "w") as config_file:
+                config_file.write(OmegaConf.to_yaml(cfg, resolve=True))
+            log.info(f"Uploaded config to {config_path}")
+        except Exception as e:
+            log.warning(f"Failed to upload config to S3: {e}")
+    else:
+        # Local path - save config
+        os.makedirs(save_dir, exist_ok=True)
+        config_path = os.path.join(save_dir, "config.yaml")
+        with open(config_path, "w") as config_file:
+            config_file.write(OmegaConf.to_yaml(cfg, resolve=True))
+        log.info(f"Saved config to {config_path}")
+
+    # Extract features from all datasets
+    datasets = []
+    if datamodule.train_dataset is not None:
+        datasets.append(datamodule.train_dataset)
+    if datamodule.val_datasets:
+        datasets.extend([d for d in datamodule.val_datasets if d is not None])
+    if datamodule.test_datasets:
+        datasets.extend([d for d in datamodule.test_datasets if d is not None])
+
+    for dataset in datasets:
+        if dataset is not None:
+            log.info(f"Extracting features from dataset: {type(dataset).__name__}")
+            dataset.extract_and_save_features(
+                model.audio_encoder,
+                save_dir=save_dir,
+                extract_method=extract_method,
+                out_key=out_key,
+                hop=hop,
+                limit_n=limit_n,
+                save=save,
+                verbose=False,
+                root_path=root_path,
+                done_ids=done_ids
+            )
+
+    log.info("Feature extraction completed!")
+    
+    return {
+        "save_dir": save_dir,
+        "extract_method": extract_method,
+        "out_key": out_key,
+    }
+
+
+@hydra_main(version_base="1.3", config_path="../configs", config_name="extract_feature.yaml")
+def main(cfg: DictConfig) -> Optional[Dict[str, Any]]:
+    """Main entry point for feature extraction.
+
+    :param cfg: DictConfig configuration composed by Hydra.
+    :return: Dict with extraction metadata.
+    """
+    # handle A100 GPUs
+    if torch.cuda.is_available() and ("A100" in torch.cuda.get_device_name() or "A5000" in torch.cuda.get_device_name()):
+        torch.set_float32_matmul_precision("high")
+
+    # avoid annoying multiprocessing errors
+    torch.multiprocessing.set_sharing_strategy('file_system')
+
+    # prevent annoying warning
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+    # apply extra utilities
+    # (e.g. ask for tags if none are provided in cfg, print cfg tree, etc.)
+    extras(cfg)
+
+    # extract features
+    result = extract_features(cfg)
+
+    return result
+
+
+if __name__ == "__main__":
+    main()
+

steerable_retrieval/models/__init__.py

@@ -0,0 +1,2 @@
+"""Model definitions for steerable_retrieval."""
+

steerable_retrieval/models/base.py

@@ -0,0 +1,219 @@
+from torch import nn
+
+class BaseModule(nn.Module):
+    def __init__(self, ckpt_path = None, freeze = False):
+        super().__init__()
+        self.ckpt_path = ckpt_path
+        self.freeze = freeze
+
+
+         
+    @classmethod
+    def from_config(cls, config, device=None):
+        config.update(device=device)
+        return cls(**config)
+    
+    @classmethod
+    def from_yaml(cls, yaml_path, device=None):
+        
+        if 's3://' in yaml_path:
+            import s3fs
+            fs = s3fs.S3FileSystem()
+            with fs.open(yaml_path, "r") as file:
+                config = yaml.safe_load(file)
+        else:
+            with open(yaml_path, "r") as file:
+                config = yaml.safe_load(file)
+        
+        config = config.get('model', config)
+        config = config.get('init_args', config)
+        
+        return cls.from_config(config, device=device)
+    
+    @classmethod
+    def from_pretrained(cls, yaml_or_config, ckpt_path, device=None):
+        if isinstance(yaml_or_config, str):
+            model = cls.from_yaml(yaml_or_config, device=device)
+        else:
+            model = cls.from_config(yaml_or_config, device=device)
+            
+        if 's3://' in ckpt_path:
+            from s3torchconnector import S3Checkpoint
+            checkpoint= S3Checkpoint(region='us-east-1')
+            with checkpoint.reader(ckpt_path) as f:
+                ckpt = torch.load(f, map_location=device)
+                model.load_state_dict(ckpt['state_dict'], strict=True)
+                print(f"Model loaded from {ckpt_path}")
+        else:
+            ckpt = torch.load(ckpt_path, map_location=device)
+            model.load_state_dict(ckpt['state_dict'], strict=True)
+            print(f"Model loaded from {ckpt_path}")
+        
+        return model
+    
+
+    def configure_optimizers(self):
+        from steerable_retrieval.utils.instantiators import instantiate
+        
+        if not hasattr(self, 'optimizer'):
+            self.optimizer = None
+        if self.optimizer is None:
+            optimizer = optim.Adam(
+                self.parameters(), lr=1e-4, betas=(0.9, 0.999), eps=1e-8)
+        else:
+            # If optimizer is a config dict or partial, instantiate it with model parameters
+            if isinstance(self.optimizer, dict):
+                # Handle _partial_ configs - Hydra creates a functools.partial
+                optimizer_cfg = dict(self.optimizer)
+                
+                # If using _target_ style (with or without _partial_)
+                if '_target_' in optimizer_cfg:
+                    # Remove _partial_ flag if present (it was just to prevent instantiation)
+                    optimizer_cfg.pop('_partial_', None)
+                    optimizer_cfg.pop('_convert_', None)
+                    optimizer_cfg['params'] = self.parameters()
+                    optimizer = instantiate(optimizer_cfg)
+                # If using class_path style
+                elif 'class_path' in optimizer_cfg:
+                    import importlib
+                    module_path, class_name = optimizer_cfg['class_path'].rsplit('.', 1)
+                    module = importlib.import_module(module_path)
+                    optimizer_class = getattr(module, class_name)
+                    init_args = optimizer_cfg.get('init_args', {})
+                    init_args['params'] = self.parameters()
+                    optimizer = optimizer_class(**init_args)
+                else:
+                    # Fallback: assume it's a direct config
+                    optimizer_cfg['params'] = self.parameters()
+                    optimizer = instantiate(optimizer_cfg)
+            elif hasattr(self.optimizer, 'func') and hasattr(self.optimizer, 'keywords'):
+                # It's a functools.partial (from _partial_=true) - call it with params
+                optimizer = self.optimizer(params=self.parameters())
+            elif callable(self.optimizer):
+                # If it's a callable (old style), call it with parameters
+                optimizer = self.optimizer(self.parameters())
+            else:
+                # Fallback to default
+                optimizer = optim.Adam(
+                    self.parameters(), lr=1e-4, betas=(0.9, 0.999), eps=1e-8)
+            
+        if hasattr(self, 'scheduler') and self.scheduler is not None:
+            # copy of the scheduler applied to the optimizer
+            ## retrocompatibilty with old schedulers
+            if isinstance(self.scheduler, dict) and 'class_name' in self.scheduler.keys():
+                scheduler_class = eval(self.scheduler['class_name'])
+                scheduler_kwargs = self.scheduler.get('init_args', {})
+                scheduler = scheduler_class(optimizer, **scheduler_kwargs)
+                self.scheduler = scheduler  # Store instantiated scheduler
+                # Return with proper Lightning configuration
+                return {
+                    'optimizer': optimizer,
+                    'lr_scheduler': {
+                        'scheduler': scheduler,
+                        'interval': 'step',  # Step after optimizer steps (respects accumulate_grad_batches)
+                        'frequency': 1,      # Step every optimizer step
+                    }
+                }
+            elif isinstance(self.scheduler, dict):
+                # Handle configs that were prevented from instantiation
+                scheduler_cfg = dict(self.scheduler)
+                
+                # If using class_path style (not _target_)
+                if 'class_path' in scheduler_cfg:
+                    import importlib
+                    module_path, class_name = scheduler_cfg['class_path'].rsplit('.', 1)
+                    module = importlib.import_module(module_path)
+                    scheduler_class = getattr(module, class_name)
+                    init_args = scheduler_cfg.get('init_args', {})
+                    init_args['optimizer'] = optimizer
+                    scheduler = scheduler_class(**init_args)
+                # If using _target_ style
+                elif '_target_' in scheduler_cfg:
+                    # Remove _partial_ flag and add optimizer
+                    scheduler_cfg.pop('_partial_', None)
+                    scheduler_cfg.pop('_convert_', None)
+                    scheduler_cfg['optimizer'] = optimizer
+                    scheduler = instantiate(scheduler_cfg)
+                else:
+                    # Fallback: assume it's a direct config
+                    scheduler_cfg['optimizer'] = optimizer
+                    scheduler = instantiate(scheduler_cfg)
+                self.scheduler = scheduler  # Store instantiated scheduler
+                # Return with proper Lightning configuration
+                return {
+                    'optimizer': optimizer,
+                    'lr_scheduler': {
+                        'scheduler': scheduler,
+                        'interval': 'step',  # Step after optimizer steps (respects accumulate_grad_batches)
+                        'frequency': 1,      # Step every optimizer step
+                    }
+                }
+            elif hasattr(self.scheduler, 'func') and hasattr(self.scheduler, 'keywords'):
+                # It's a functools.partial (from _partial_=true) - call it with optimizer
+                scheduler = self.scheduler(optimizer=optimizer)
+                self.scheduler = scheduler  # Store instantiated scheduler
+                # Return with proper Lightning configuration
+                return {
+                    'optimizer': optimizer,
+                    'lr_scheduler': {
+                        'scheduler': scheduler,
+                        'interval': 'step',  # Step after optimizer steps (respects accumulate_grad_batches)
+                        'frequency': 1,      # Step every optimizer step
+                    }
+                }
+            else:
+                # If scheduler is already instantiated, just return it
+                # Return with proper Lightning configuration
+                return {
+                    'optimizer': optimizer,
+                    'lr_scheduler': {
+                        'scheduler': self.scheduler,
+                        'interval': 'step',  # Step after optimizer steps (respects accumulate_grad_batches)
+                        'frequency': 1,      # Step every optimizer step
+                    }
+                }
+        
+        return optimizer
+
+    def load_ckpt(self, ckpt_path, device = None, prefix = ''):
+
+        if device is None:
+            device = next(self.parameters()).device
+            
+        if 's3://' in ckpt_path:
+            from s3torchconnector import S3Checkpoint
+            checkpoint= S3Checkpoint(region='us-east-1')
+            with checkpoint.reader(ckpt_path) as f:
+                state_dict = torch.load(f, map_location=device)['state_dict']
+                print(f"Model loaded from {ckpt_path}")
+        else:
+            state_dict = torch.load(ckpt_path, map_location=device)['state_dict']
+            print(f"Model loaded from {ckpt_path}")
+        
+        print(state_dict)
+        
+        try:
+            self.load_state_dict(state_dict)
+            print("Loaded full state dict")
+        except:
+            print("Could not load state dict, trying to load only ['encoder'] keys")
+            
+            try:
+                from collections import OrderedDict
+                new_state_dict = OrderedDict()
+                for k in list(state_dict.keys()):
+                    if prefix in k:
+                        new_key = k.replace('encoder.','')
+                        new_state_dict[new_key] = state_dict[k]
+                        
+                self.load_state_dict(new_state_dict)
+                print(f"Loaded only {prefix} keys")
+                
+            except Exception as e:
+                print(f"Could not load state dict, error: {e}")
+            
+        
+
+    def freeze(self):
+        for param in self.parameters():
+            param.requires_grad = False

steerable_retrieval/models/encoders/__init__.py

@@ -0,0 +1,11 @@
+"""Audio and text encoder wrappers."""
+
+from steerable_retrieval.models.encoders.clap import CLAPAudioEncoder, CLAPTextEncoder
+from steerable_retrieval.models.encoders.muq import MuQAudioEncoder, MuQTextEncoder
+
+__all__ = [
+    'CLAPAudioEncoder',
+    'CLAPTextEncoder',
+    'MuQAudioEncoder',
+    'MuQTextEncoder',
+]

steerable_retrieval/models/encoders/clap.py

@@ -0,0 +1,158 @@
+"""CLAP encoder wrappers for audio and text."""
+
+import torch
+import torch.nn as nn
+from transformers import ClapModel, ClapProcessor
+from steerable_retrieval.models.base import BaseModule
+import numpy as np
+
+def _to_embedding_tensor(out):
+    """Return [batch, feature_dim] tensor from model output (tensor or HuggingFace output)."""
+    if isinstance(out, torch.Tensor):
+        return out
+    if hasattr(out, "pooler_output") and out.pooler_output is not None:
+        return out.pooler_output
+    if hasattr(out, "last_hidden_state"):
+        h = out.last_hidden_state
+        return h[:, 0] if h.dim() == 3 else h
+    raise TypeError(
+        f"Model returned {type(out).__name__}; expected tensor or pooler_output/last_hidden_state."
+    )
+
+import logging
+
+
+class CLAPAudioEncoder(BaseModule):
+    """CLAP audio encoder wrapper that returns dense features after projection head."""
+
+    def __init__(self, model_name="laion/larger_clap_music", device="cuda", freeze=True, ckpt_path=None, sampling_rate=48000, **kwargs):
+        super().__init__(ckpt_path=ckpt_path, freeze=freeze)
+        self.model_name = model_name
+        self.freeze = freeze
+        self.sampling_rate = sampling_rate
+
+        # Load CLAP model and processor
+        self.model = ClapModel.from_pretrained(model_name)
+        self.processor = ClapProcessor.from_pretrained(model_name)
+        # Set sampling rate to suppress warning
+        if hasattr(self.processor, 'feature_extractor'):
+            self.processor.feature_extractor.sampling_rate = sampling_rate
+
+        # Set to eval mode and freeze if requested
+        self.model.eval()
+        if freeze:
+            for param in self.model.parameters():
+                param.requires_grad = False
+
+        
+
+    def forward(self, audio):
+        """
+        Forward pass through CLAP audio encoder.
+        
+        Args:
+            audio: Audio input. Can be:
+                - Raw audio array (numpy array or torch tensor)
+                - Dictionary with 'array' key (from datasets)
+                - Already processed audio tensor
+        
+        Returns:
+            audio_features: Dense audio features [batch, feature_dim]
+        """
+        # Process audio if needed
+        device = next(self.model.parameters()).device
+        if isinstance(audio, dict) and 'array' in audio:
+            # Handle dataset format
+            audio_array = audio['array']
+            inputs = self.processor(audio=audio_array, return_tensors="pt", sampling_rate=self.sampling_rate).to(device)
+        elif isinstance(audio, torch.Tensor):
+            # Handle batched audio tensors from DataLoader
+            # DataLoader stacks tensors, so we get [batch, 1, samples] or [batch, samples]
+            if audio.dim() == 3:
+                # [batch, 1, samples] -> squeeze middle dimension
+                audio = audio.squeeze(1)
+            elif audio.dim() == 2:
+                # [batch, samples] - already correct
+                pass
+            elif audio.dim() == 1:
+                # [samples] - single sample, add batch dim
+                audio = audio.unsqueeze(0)
+            
+            # Convert to list of numpy arrays for CLAP processor
+            # CLAP processor expects list of 1D arrays
+            audio_list = [a.cpu().numpy() for a in audio]
+            inputs = self.processor(audio=audio_list, return_tensors="pt", sampling_rate=self.sampling_rate).to(device)
+        else:
+            # Assume it's a numpy array or list
+            inputs = self.processor(audio=audio, return_tensors="pt", sampling_rate=self.sampling_rate).to(device)
+
+        
+        # make sure the inputs are on the correct device
+        for key, value in inputs.items():
+            # if numpy array, convert to tensor
+            if isinstance(value, np.ndarray):
+                value = torch.tensor(value).to(device)
+                inputs[key] = value
+
+        # Get dense audio features (after projection head)
+        if self.freeze:
+            with torch.no_grad():
+                out = self.model.get_audio_features(**inputs)
+        else:
+            out = self.model.get_audio_features(**inputs)
+        return _to_embedding_tensor(out)
+
+
+class CLAPTextEncoder(BaseModule):
+    """CLAP text encoder wrapper that returns dense features after projection head."""
+
+    def __init__(self, model_name="laion/larger_clap_music", device="cuda", freeze=True, ckpt_path=None, **kwargs):
+        super().__init__(ckpt_path=ckpt_path, freeze=freeze)
+        self.model_name = model_name
+        self.device = device
+        self.freeze = freeze
+
+        # Load CLAP model and processor
+        self.model = ClapModel.from_pretrained(model_name).to(device)
+        self.processor = ClapProcessor.from_pretrained(model_name)
+
+        # Set to eval mode and freeze if requested
+        self.model.eval()
+        if freeze:
+            for param in self.model.parameters():
+                param.requires_grad = False
+
+    def forward(self, text):
+        """
+        Forward pass through CLAP text encoder.
+        
+        Args:
+            text: Text input. Can be:
+                - String
+                - List of strings
+                - Already processed text
+        
+        Returns:
+            text_features: Dense text features [batch, feature_dim]
+        """
+        # Process text if needed
+        if isinstance(text, str):
+            text = [text]
+
+        device = next(self.model.parameters()).device
+        
+        # Process text through CLAP processor
+        inputs = self.processor(text=text, return_tensors="pt", padding=True, truncation=True).to(self.device)
+
+        # make sure the inputs are on the correct device
+        for key, value in inputs.items():
+            # if numpy array, convert to tensor
+            if isinstance(value, np.ndarray):
+                value = torch.tensor(value).to(device)
+                inputs[key] = value
+                
+
+        # Get dense text features (after projection head)
+        with torch.set_grad_enabled(not self.freeze):
+            out = self.model.get_text_features(**inputs)
+        return _to_embedding_tensor(out)

steerable_retrieval/models/encoders/muq.py

@@ -0,0 +1,120 @@
+"""MuQ encoder wrappers for audio and text."""
+
+import torch
+import torch.nn as nn
+from muq import MuQMuLan
+from steerable_retrieval.models.base import BaseModule
+
+
+class MuQAudioEncoder(BaseModule):
+    """MuQ audio encoder wrapper that returns dense features after projection head."""
+
+    def __init__(self, model_name="OpenMuQ/MuQ-MuLan-large", device="cuda", freeze=True, ckpt_path=None, sampling_rate=24000, **kwargs):
+        super().__init__(ckpt_path=ckpt_path, freeze=freeze)
+        self.model_name = model_name
+        self.device = device
+        self.freeze = freeze
+
+        # Load MuQ model
+        self.model = MuQMuLan.from_pretrained(model_name)
+        self.model = self.model.to(device).eval()
+
+        # Freeze if requested
+        if freeze:
+            for param in self.model.parameters():
+                param.requires_grad = False
+
+        self.sampling_rate = sampling_rate
+
+    def forward(self, audio):
+        """
+        Forward pass through MuQ audio encoder.
+        
+        Args:
+            audio: Audio input. Can be:
+                - Raw audio tensor [batch, samples] or [samples]
+                - Numpy array
+                - Already processed audio
+        
+        Returns:
+            audio_features: Dense audio features [batch, feature_dim]
+        """
+        # Ensure audio is a tensor on the correct device
+        if not isinstance(audio, torch.Tensor):
+            audio = torch.tensor(audio)
+        
+        # Ensure audio is on the correct device
+        audio = audio.to(self.device)
+        
+        # Handle batched audio tensors from DataLoader
+        # DataLoader stacks tensors, so we get [batch, 1, samples] or [batch, samples]
+        if audio.dim() == 3:
+            # [batch, 1, samples] -> squeeze middle dimension
+            audio = audio.squeeze(1)
+        elif audio.dim() == 2:
+            # [batch, samples] - already correct
+            pass
+        elif audio.dim() == 1:
+            # [samples] - single sample, add batch dim
+            audio = audio.unsqueeze(0)
+        
+        # Get dense audio features (after projection head)
+        with torch.set_grad_enabled(not self.freeze):
+            out = self.model(wavs=audio)
+        # MuQ/HuggingFace models may return BaseModelOutputWithPooling; return tensor [B, D]
+        if isinstance(out, torch.Tensor):
+            return out
+        if hasattr(out, "pooler_output") and out.pooler_output is not None:
+            return out.pooler_output
+        if hasattr(out, "last_hidden_state"):
+            h = out.last_hidden_state
+            return h[:, 0] if h.dim() == 3 else h
+        raise TypeError(f"MuQ model returned {type(out).__name__}; expected tensor or pooler_output/last_hidden_state.")
+
+
+class MuQTextEncoder(BaseModule):
+    """MuQ text encoder wrapper that returns dense features after projection head."""
+
+    def __init__(self, model_name="OpenMuQ/MuQ-MuLan-large", device="cuda", freeze=True, ckpt_path=None, **kwargs):
+        super().__init__(ckpt_path=ckpt_path, freeze=freeze)
+        self.model_name = model_name
+        self.device = device
+        self.freeze = freeze
+
+        # Load MuQ model
+        self.model = MuQMuLan.from_pretrained(model_name)
+        self.model = self.model.to(device).eval()
+
+        # Freeze if requested
+        if freeze:
+            for param in self.model.parameters():
+                param.requires_grad = False
+
+    def forward(self, text):
+        """
+        Forward pass through MuQ text encoder.
+        
+        Args:
+            text: Text input. Can be:
+                - String
+                - List of strings
+        
+        Returns:
+            text_features: Dense text features [batch, feature_dim]
+        """
+        # Ensure text is a list
+        if isinstance(text, str):
+            text = [text]
+        
+        # Get dense text features (after projection head)
+        with torch.set_grad_enabled(not self.freeze):
+            out = self.model(texts=text)
+        # MuQ/HuggingFace models may return BaseModelOutputWithPooling; return tensor [B, D]
+        if isinstance(out, torch.Tensor):
+            return out
+        if hasattr(out, "pooler_output") and out.pooler_output is not None:
+            return out.pooler_output
+        if hasattr(out, "last_hidden_state"):
+            h = out.last_hidden_state
+            return h[:, 0] if h.dim() == 3 else h
+        raise TypeError(f"MuQ model returned {type(out).__name__}; expected tensor or pooler_output/last_hidden_state.")

steerable_retrieval/models/sae/__init__.py

@@ -0,0 +1,46 @@
+"""Sparse Autoencoder (SAE) implementations."""
+
+from steerable_retrieval.models.sae.sae import SAE, LightningSAE
+from steerable_retrieval.models.sae.encoders import (
+    SAEEncoder,
+    VanillaSAEEncoder,
+    TopKSAEEncoder,
+    BatchTopKSAEEncoder,
+    MatryoshkaBatchTopKSAEEncoder,
+    JumpReLUSAEEncoder,
+    OMPEncoder,
+    MPEncoder,
+)
+from steerable_retrieval.models.sae.decoders import SAEDecoder
+from steerable_retrieval.models.sae.penalties import (
+    VanillaPenalty,
+    TopKPenalty,
+    BatchTopKPenalty,
+    MatryoshkaBatchTopKPenalty,
+    L1Penalty,
+    JumpReLU,
+    StepFunction,
+    RectangleFunction,
+)
+
+__all__ = [
+    'SAE',
+    'LightningSAE',
+    'SAEEncoder',
+    'VanillaSAEEncoder',
+    'TopKSAEEncoder',
+    'BatchTopKSAEEncoder',
+    'MatryoshkaBatchTopKSAEEncoder',
+    'JumpReLUSAEEncoder',
+    'OMPEncoder',
+    'MPEncoder',
+    'SAEDecoder',
+    'VanillaPenalty',
+    'TopKPenalty',
+    'BatchTopKPenalty',
+    'MatryoshkaBatchTopKPenalty',
+    'L1Penalty',
+    'JumpReLU',
+    'StepFunction',
+    'RectangleFunction',
+]

steerable_retrieval/models/sae/decoders.py

@@ -0,0 +1,83 @@
+"""SAE decoder classes that return only reconstruction."""
+
+import torch
+import torch.nn as nn
+
+
+class SAEDecoder(nn.Module):
+    """Base SAE decoder that takes activations and returns reconstruction."""
+
+    def __init__(self, dict_size, act_size, device='cuda', 
+                 input_unit_norm=False, seed=None, output_unit_norm=False):
+        """
+        Initialize decoder.
+        
+        Args:
+            dict_size: Dictionary size (number of features)
+            act_size: Activation size (input/output feature dimension)
+            device: Device to use
+            input_unit_norm: Whether to apply unit norm postprocessing
+            seed: Random seed for initialization
+        """
+        super().__init__()
+        self.dict_size = dict_size
+        self.act_size = act_size
+        self.device = device
+        self.dtype = torch.float32  # Automatic dtype
+        self.input_unit_norm = input_unit_norm
+        self.output_unit_norm = output_unit_norm
+        if seed is not None:
+            torch.manual_seed(seed)
+        
+        # Decoder weights - initialize randomly, can be set from encoder later
+        self.W_dec = nn.Parameter(
+            torch.nn.init.kaiming_uniform_(
+                torch.empty(dict_size, act_size)
+            )
+        )
+        
+        # b_dec will be set separately to share with encoder
+        # Set it after initialization: decoder.b_dec = encoder.b_dec
+        
+        self.to(self.dtype).to(device)
+        
+        # For compatibility
+        self.in_channels = dict_size
+        self.out_channels = act_size
+
+    def forward(self, activations):
+        """
+        Forward pass through SAE decoder.
+        
+        Args:
+            activations: Feature activations [batch, ..., dict_size]
+            
+        Returns:
+            reconstruction: Reconstructed features [batch, ..., act_size]
+        """
+        if not hasattr(self, 'b_dec'):
+            raise RuntimeError("b_dec not set. This should be set automatically by SAE/LightningSAE during initialization.")
+        
+        reconstruction = activations @ self.W_dec + self.b_dec
+        
+        # Apply postprocessing if needed
+        if self.input_unit_norm:
+            # Note: We don't have x_mean and x_std here, so we skip postprocessing
+            # This is fine since preprocessing/postprocessing are typically used together
+            pass
+        
+        if self.output_unit_norm:
+            reconstruction = reconstruction / reconstruction.norm(dim=-1, keepdim=True)
+        
+        return reconstruction
+
+    @torch.no_grad()
+    def make_decoder_weights_and_grad_unit_norm(self):
+        """Normalize decoder weights and adjust gradients."""
+        W_dec_normed = self.W_dec / self.W_dec.norm(dim=-1, keepdim=True)
+        if self.W_dec.grad is not None:
+            W_dec_grad_proj = (self.W_dec.grad * W_dec_normed).sum(
+                -1, keepdim=True
+            ) * W_dec_normed
+            self.W_dec.grad -= W_dec_grad_proj
+        self.W_dec.data = W_dec_normed

steerable_retrieval/models/sae/encoders.py

@@ -0,0 +1,360 @@
+"""SAE encoder classes that return only activations."""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+try:
+    from overcomplete.optimization import batched_matrix_nnls
+except ImportError:
+    batched_matrix_nnls = None
+
+from steerable_retrieval.models.sae.penalties import (
+    VanillaPenalty, TopKPenalty, BatchTopKPenalty,
+    MatryoshkaBatchTopKPenalty, JumpReLU,
+)
+
+
+class SAEEncoder(nn.Module):
+    """Base SAE encoder that takes audio/text and returns activations."""
+
+    def __init__(self, act_size, dict_size, device='cuda', 
+                 input_unit_norm=False, seed=None, concat_dim=-1, **kwargs):
+        super().__init__()
+        self.act_size = act_size
+        self.dict_size = dict_size
+        self.device = device
+        self.dtype = torch.float32  # Automatic dtype
+        self.input_unit_norm = input_unit_norm
+        self.concat_dim = concat_dim
+        self.cfg = {
+            "act_size": act_size,
+            "dict_size": dict_size,
+            "device": device,
+            "dtype": self.dtype,
+            "input_unit_norm": input_unit_norm,
+            **kwargs
+        }
+        
+        if seed is not None:
+            torch.manual_seed(seed)
+
+        # Encoder weights and biases
+        self.b_dec = nn.Parameter(torch.zeros(act_size))
+        self.b_enc = nn.Parameter(torch.zeros(dict_size))
+        self.W_enc = nn.Parameter(
+            torch.nn.init.kaiming_uniform_(
+                torch.empty(act_size, dict_size)
+            )
+        )
+        
+        # Track inactive features
+        self.num_batches_not_active = torch.zeros((dict_size,)).to(device)
+
+        self.to(self.dtype).to(device)
+        
+        # Required by SAE class
+        self.out_channels = dict_size
+
+    def forward(self, x):
+        """
+        Forward pass through SAE encoder.
+        
+        Args:
+            x: Input features tensor [batch, ..., act_size]
+               Can be audio or text features (already encoded by audio_encoder/text_encoder)
+            
+        Returns:
+            activations: Feature activations [batch, ..., dict_size]
+        """
+        # Encode to get activations
+        pre, z = self._encode(x)
+
+        self.update_inactive_features(z)
+        
+        return pre, z
+
+    def _encode(self, x):
+        """Internal encode method - to be overridden by subclasses."""
+        raise NotImplementedError
+
+    def preprocess_input(self, x):
+        if self.input_unit_norm:
+            x_mean = x.mean(dim=-1, keepdim=True)
+            x = x - x_mean
+            # unit norm the input
+            x = x / (x.norm(dim=-1, keepdim=True) + 1e-6)
+            return x, x_mean, None
+        else:
+            return x, None, None
+
+    def update_inactive_features(self, acts):
+        self.num_batches_not_active += (acts.sum(0) == 0).float()
+        self.num_batches_not_active[acts.sum(0) > 0] = 0
+
+
+class VanillaSAEEncoder(SAEEncoder):
+    """Vanilla SAE encoder with standard ReLU activation."""
+
+    def _encode(self, x):
+        x, x_mean, x_std = self.preprocess_input(x)
+        x_cent = x - self.b_dec
+        pre_activations = x_cent @ self.W_enc + self.b_enc
+        acts = F.relu(pre_activations)
+        return pre_activations, acts
+
+
+class TopKSAEEncoder(SAEEncoder):
+    """TopK SAE encoder - keeps top k activations per sample."""
+
+    def __init__(self, act_size, dict_size, top_k, device='cuda', 
+                 input_unit_norm=False, seed=None, concat_dim=-1, **kwargs):
+        super().__init__(act_size, dict_size, device, input_unit_norm, seed, concat_dim, top_k=top_k, **kwargs)
+        self.top_k = top_k
+        self.penalty = TopKPenalty(top_k)
+
+    def _init_decoder_weights(self, W_dec):
+        # initialize the encoder weight to the transpose of the decoder weight
+        self.W_enc.data.copy_(W_dec.T)
+
+    def _encode(self, x):
+        x, _, _ = self.preprocess_input(x)
+        x_cent = x - self.b_dec
+        pre_activations = x_cent @ self.W_enc
+        acts = F.relu(pre_activations)
+        acts = self.penalty(acts)
+        return pre_activations, acts
+
+
+class BatchTopKSAEEncoder(SAEEncoder):
+    """
+    Batch Top-k SAE encoder (Bussmann et al., 2024).
+
+    Retains only the top-k global activations across the entire batch.
+    During training, the k-th highest activation (over the flattened batch)
+    is used as the threshold and a running average is maintained with
+    exponential momentum.  At eval time the running threshold is used,
+    making inference deterministic and batch-size independent.
+    """
+
+    def __init__(self, act_size, dict_size, top_k, threshold_momentum=0.9,
+                 device='cuda', input_unit_norm=False, seed=None, concat_dim=-1, **kwargs):
+        super().__init__(act_size, dict_size, device, input_unit_norm, seed, concat_dim,
+                         top_k=top_k, threshold_momentum=threshold_momentum, **kwargs)
+        self.top_k = top_k
+        self.penalty = BatchTopKPenalty(top_k, threshold_momentum=threshold_momentum)
+        self.register_buffer("running_threshold", None)
+        self.threshold_momentum = threshold_momentum
+    def _init_decoder_weights(self, W_dec):
+        self.W_enc.data.copy_(W_dec.T)
+
+    def _encode(self, x):
+        x, _, _ = self.preprocess_input(x)
+        x_cent = x - self.b_dec
+        pre_activations = x_cent @ self.W_enc
+        acts = F.relu(pre_activations)
+        threshold = self.penalty._get_threshold(acts)
+        threshold = self._update_threshold(threshold)
+        acts = self.penalty(acts, threshold)
+        return pre_activations, acts
+
+    def _rectify(self, acts):
+        threshold = self.penalty._get_threshold(acts)
+        threshold = self._update_threshold(threshold)
+        acts = self.penalty(acts, threshold)
+        return acts
+    
+    def _update_threshold(self, threshold):
+        if self.training:
+            if self.running_threshold is None:
+                self.running_threshold = threshold.detach()
+            else:
+                self.running_threshold = (
+                    self.threshold_momentum * self.running_threshold
+                    + (1 - self.threshold_momentum) * threshold.detach()
+                )
+            return threshold
+        else:
+            if self.running_threshold is None:
+                self.running_threshold = threshold.detach()
+                return threshold
+            else:
+                return self.running_threshold
+    
+
+    
+
+
+class MatryoshkaBatchTopKSAEEncoder(BatchTopKSAEEncoder):
+    """
+    Matryoshka Batch Top-k SAE encoder.
+
+    Combines nested feature groups (Matryoshka) with global batch-level TopK
+    sparsity.  ``group_sizes`` defines the nested structure and implicitly sets
+    ``dict_size = sum(group_sizes)``.  Features beyond the currently active
+    groups are zeroed out by the penalty.
+    """
+
+    def __init__(self, act_size, group_sizes, top_k, threshold_momentum=0.9,
+                 device='cuda', input_unit_norm=False, seed=None, concat_dim=-1, **kwargs):
+        dict_size = sum(group_sizes)
+        super().__init__(act_size, dict_size, top_k, threshold_momentum,
+                         device, input_unit_norm, seed, concat_dim, **kwargs)
+        self.group_sizes = group_sizes
+        # Replace the BatchTopKPenalty created by parent with Matryoshka variant
+        self.penalty = MatryoshkaBatchTopKPenalty(top_k, group_sizes, threshold_momentum)
+
+
+class JumpReLUSAEEncoder(SAEEncoder):
+    """JumpReLU SAE encoder with learnable thresholds."""
+
+    def __init__(self, act_size, dict_size, bandwidth=0.1, device='cuda', 
+                 input_unit_norm=False, use_pre_enc_bias=False, seed=None, concat_dim=-1, **kwargs):
+        super().__init__(act_size, dict_size, device, input_unit_norm, seed, concat_dim, 
+                        bandwidth=bandwidth, use_pre_enc_bias=use_pre_enc_bias, **kwargs)
+        self.bandwidth = bandwidth
+        self.use_pre_enc_bias = use_pre_enc_bias
+        self.jumprelu = JumpReLU(feature_size=dict_size, bandwidth=bandwidth, device=device)
+
+    def _encode(self, x):
+        x, _, _ = self.preprocess_input(x)
+        
+        if self.use_pre_enc_bias:
+            x = x - self.b_dec
+
+        pre_activations = x @ self.W_enc + self.b_enc
+        acts = self.jumprelu(acts)
+        return pre_activations, acts
+
+
+class OMPEncoder(SAEEncoder):
+    """
+    Orthogonal Matching Pursuit SAE encoder.
+
+    Uses OMP to find sparse codes: at each iteration, select the atom most correlated
+    with the residual, then solve NNLS over selected atoms. Requires the decoder to be
+    set via set_decoder(decoder) so the encoder can use decoder.W_dec as the dictionary.
+    Encoding is non-differentiable (no_grad). Use penalty=None when using this encoder.
+    """
+
+    def __init__(self, act_size, dict_size, k=1, dropout=None, max_iter=10,
+                 device='cuda', input_unit_norm=False, seed=None, concat_dim=-1, **kwargs):
+        assert isinstance(k, int) and k > 0, "k must be a positive integer."
+        if dropout is not None:
+            assert 0.0 <= dropout <= 1.0, "Dropout must be in range [0, 1]."
+        assert isinstance(max_iter, int) and max_iter > 0, "max_iter must be a positive integer."
+        if batched_matrix_nnls is None:
+            raise ImportError("OMPEncoder requires overcomplete.optimization.batched_matrix_nnls")
+        super().__init__(act_size, dict_size, device, input_unit_norm, seed, concat_dim, **kwargs)
+        self.k = k
+        self.dropout = dropout
+        self.max_iter = max_iter
+        self._decoder = None
+
+    def set_decoder(self, decoder):
+        """Set the decoder so this encoder can use its W_dec as the dictionary."""
+        self._decoder = decoder
+
+    def _encode(self, x):
+        if self._decoder is None:
+            raise RuntimeError("OMPEncoder requires decoder to be set. Call set_decoder(decoder) or use SAE/LightningSAE which does this automatically.")
+        W = self._decoder.W_dec  # [dict_size, act_size]
+        device = W.device
+        dtype = W.dtype
+        k = self.k
+        max_iter = self.max_iter
+
+        if self.dropout is not None and self.training:
+            drop_mask = torch.bernoulli(
+                (1.0 - self.dropout) * torch.ones(W.shape[0], device=device, dtype=dtype)
+            )
+            W = W * drop_mask.unsqueeze(1)
+
+        x, _, _ = self.preprocess_input(x)
+        x = x.to(device=device, dtype=dtype)
+        batch_size = x.shape[0]
+        codes = torch.zeros(batch_size, self.dict_size, device=device, dtype=dtype)
+        residual = x.clone()
+        selected_atoms = None
+
+        with torch.no_grad():
+            for _ in range(k):
+                z_corr = residual @ W.T  # [B, dict_size]
+                if selected_atoms is not None:
+                    z_corr.scatter_(dim=1, index=selected_atoms, value=-torch.inf)
+                _, idx = torch.topk(z_corr, k=1, dim=1)  # [B, 1]
+                selected_atoms = idx if selected_atoms is None else torch.cat([selected_atoms, idx], dim=1)  # [B, num_selected]
+                W_sel = W[selected_atoms]  # [B, num_selected, act_size]
+                Z_init = torch.gather(codes, 1, selected_atoms)  # [B, num_selected]
+                codes_sel = batched_matrix_nnls(
+                    W_sel, x, max_iter=max_iter, tol=1e-5, Z_init=Z_init
+                )
+                codes.scatter_(dim=1, index=selected_atoms, src=codes_sel)
+                residual = x - codes @ W
+
+        return None, codes
+
+    def train(self, mode=True):
+        if not mode:
+            self.dropout = None
+        return super().train(mode)
+
+
+class MPEncoder(SAEEncoder):
+    """
+    Matching Pursuit SAE encoder.
+
+    Greedy MP: at each of k iterations, pick the atom most correlated with the
+    residual, add that contribution to the codes, and subtract it from the residual.
+    No NNLS (unlike OMP). Requires the decoder to be set via set_decoder(decoder).
+    Encoding is non-differentiable (no_grad). Use penalty=None when using this encoder.
+    """
+
+    def __init__(self, act_size, dict_size, k=1, dropout=None,
+                 device='cuda', input_unit_norm=False, seed=None, concat_dim=-1, **kwargs):
+        assert isinstance(k, int) and k > 0, "k must be a positive integer."
+        if dropout is not None:
+            assert 0.0 <= dropout <= 1.0, "Dropout must be in range [0, 1]."
+        super().__init__(act_size, dict_size, device, input_unit_norm, seed, concat_dim, **kwargs)
+        self.k = k
+        self.dropout = dropout
+        self._decoder = None
+
+    def set_decoder(self, decoder):
+        """Set the decoder so this encoder can use its W_dec as the dictionary."""
+        self._decoder = decoder
+
+    def _encode(self, x):
+        if self._decoder is None:
+            raise RuntimeError(
+                "MPEncoder requires decoder to be set. Call set_decoder(decoder) or use SAE/LightningSAE which does this automatically."
+            )
+        W = self._decoder.W_dec  # [dict_size, act_size]
+        device = W.device
+        dtype = W.dtype
+
+        if self.dropout is not None and self.training:
+            drop_w = torch.bernoulli(
+                (1.0 - self.dropout) * torch.ones(W.shape[0], 1, device=device, dtype=dtype)
+            )
+            W = W * drop_w
+
+        x, _, _ = self.preprocess_input(x)
+        x = x.to(device=device, dtype=dtype)
+        codes = torch.zeros(x.shape[0], self.dict_size, device=device, dtype=dtype)
+        residual = x.clone()
+
+        with torch.no_grad():
+            for _ in range(self.k):
+                z = residual @ W.T  # [B, dict_size]
+                val, idx = torch.max(z, dim=1)  # val [B], idx [B]
+                to_add = F.one_hot(idx, num_classes=self.dict_size).float() * val.unsqueeze(1)
+                codes = codes + to_add
+                residual = residual - to_add @ W
+
+        return None, codes
+
+    def train(self, mode=True):
+        if not mode:
+            self.dropout = None
+        return super().train(mode)

steerable_retrieval/models/sae/penalties.py

@@ -0,0 +1,152 @@
+"""Penalty functions for SAE activations."""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import torch.autograd as autograd
+
+
+class VanillaPenalty(nn.Module):
+    """No penalty - just pass through activations."""
+    
+    def forward(self, activations):
+        return activations
+
+
+class TopKPenalty(nn.Module):
+    """TopK penalty - keep only top k activations per sample."""
+    
+    def __init__(self, top_k):
+        super().__init__()
+        self.top_k = top_k
+    
+    def forward(self, activations):
+        acts_topk = torch.topk(activations, self.top_k, dim=-1)
+        acts_topk = torch.zeros_like(activations).scatter(
+            -1, acts_topk.indices, acts_topk.values
+        )
+        return acts_topk
+
+
+class BatchTopKPenalty(nn.Module):
+    """
+    Batch Top-k penalty (Bussmann et al., 2024).
+
+    During training, finds the k-th highest activation across the flattened
+    batch and uses it as a threshold.  A running average of this threshold is
+    maintained via exponential momentum so that at eval time the threshold is
+    deterministic and batch-size independent.
+    """
+
+    def __init__(self, top_k, threshold_momentum=0.9):
+        super().__init__()
+        self.top_k = top_k
+        self.threshold_momentum = threshold_momentum
+
+    def forward(self, activations, threshold=None):
+        mask = (activations >= threshold).float().detach()
+        return activations * mask
+    
+    def _get_threshold(self, activations):
+        acts_topk = torch.topk(activations.flatten(), self.top_k * activations.shape[0], dim = -1)
+        return acts_topk.values[-1]
+
+
+
+class MatryoshkaBatchTopKPenalty(BatchTopKPenalty):
+    """
+    Batch Top-k penalty with Matryoshka group masking.
+
+    After applying the global BatchTopK threshold, features beyond the
+    currently active groups are zeroed out.  ``active_groups`` defaults to
+    all groups and can be reduced at inference time to use fewer features.
+    """
+
+    def __init__(self, top_k, group_sizes, threshold_momentum=0.9):
+        super().__init__(top_k, threshold_momentum)
+        self.group_sizes = group_sizes
+        self.group_indices = [0] + torch.cumsum(torch.tensor(group_sizes), dim=0).tolist()
+        self.active_groups = len(group_sizes)
+
+    def forward(self, activations, threshold=None):
+        acts = super().forward(activations, threshold)
+        max_idx = self.group_indices[self.active_groups]
+        if max_idx < acts.shape[-1]:
+            acts = acts.clone()
+            acts[..., max_idx:] = 0
+        return acts
+
+
+class L1Penalty(nn.Module):
+    """L1 penalty - ReLU activation."""
+    
+    def forward(self, activations):
+        return F.relu(activations)
+
+
+
+class RectangleFunction(autograd.Function):
+    @staticmethod
+    def forward(ctx, x):
+        ctx.save_for_backward(x)
+        return ((x > -0.5) & (x < 0.5)).float()
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (x,) = ctx.saved_tensors
+        grad_input = grad_output.clone()
+        grad_input[(x <= -0.5) | (x >= 0.5)] = 0
+        return grad_input
+
+
+class JumpReLUFunction(autograd.Function):
+    @staticmethod
+    def forward(ctx, x, log_threshold, bandwidth):
+        ctx.save_for_backward(x, log_threshold, torch.tensor(bandwidth))
+        threshold = torch.exp(log_threshold)
+        return x * (x > threshold).float()
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        x, log_threshold, bandwidth_tensor = ctx.saved_tensors
+        bandwidth = bandwidth_tensor.item()
+        threshold = torch.exp(log_threshold)
+        x_grad = (x > threshold).float() * grad_output
+        threshold_grad = (
+            -(threshold / bandwidth)
+            * RectangleFunction.apply((x - threshold) / bandwidth)
+            * grad_output
+        )
+        return x_grad, threshold_grad, None
+
+
+class JumpReLU(nn.Module):
+    def __init__(self, feature_size, bandwidth, device='cpu'):
+        super(JumpReLU, self).__init__()
+        self.log_threshold = nn.Parameter(torch.zeros(feature_size, device=device))
+        self.bandwidth = bandwidth
+
+    def forward(self, x):
+        return JumpReLUFunction.apply(x, self.log_threshold, self.bandwidth)
+
+
+class StepFunction(autograd.Function):
+    @staticmethod
+    def forward(ctx, x, log_threshold, bandwidth):
+        ctx.save_for_backward(x, log_threshold, torch.tensor(bandwidth))
+        threshold = torch.exp(log_threshold)
+        return (x > threshold).float()
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        x, log_threshold, bandwidth_tensor = ctx.saved_tensors
+        bandwidth = bandwidth_tensor.item()
+        threshold = torch.exp(log_threshold)
+        x_grad = torch.zeros_like(x)
+        threshold_grad = (
+            -(1.0 / bandwidth)
+            * RectangleFunction.apply((x - threshold) / bandwidth)
+            * grad_output
+        )
+        return x_grad, threshold_grad, None

steerable_retrieval/models/sae/sae.py

@@ -0,0 +1,346 @@
+import torch
+import torch.nn as nn
+import lightning as L
+from steerable_retrieval.models.base import BaseModule
+from steerable_retrieval.utils.pylogger import RankedLogger
+
+log = RankedLogger(__name__, rank_zero_only=True)
+
+# Loss key -> log prefix (term first: L0/Train/..., L1/val/dataset/...)
+LOSS_KEY_TO_PREFIX = {
+    "l0_norm": "L0",
+    "l1_loss": "L1",
+    "l2_loss": "L2",
+    "l1_norm": "L1_norm",
+    "aux_loss": "aux_loss",
+    "loss": "loss",
+    "mse": "mse",
+    "r2": "r2",
+    "ndn": "ndn",
+    "min_l2_loss": "min_L2",
+    "max_l2_loss": "max_L2",
+}
+
+import logging
+
+class SAE(BaseModule):
+
+    def __init__(self,
+                 audio_encoder,
+                 text_encoder,
+                 sae_encoder,
+                 sae_decoder,
+                 ckpt_path = None,
+                 freeze = False,
+                 tied = False,
+                 **kwargs):
+        super().__init__(ckpt_path=ckpt_path, freeze=freeze)
+        self.ckpt_path = ckpt_path
+        self.freeze = freeze
+
+        # Store encoder/decoder components as instance attributes
+        self.audio_encoder = audio_encoder
+        self.text_encoder = text_encoder
+        self.sae_encoder = sae_encoder
+        self.sae_decoder = sae_decoder
+
+        # Share b_dec between encoder and decoder
+        if not hasattr(self.sae_decoder, 'b_dec'):
+            self.sae_decoder.b_dec = self.sae_encoder.b_dec
+
+        # OMP-style encoders need the decoder's dictionary (W_dec) at encode time
+        if getattr(self.sae_encoder, 'set_decoder', None) is not None:
+            self.sae_encoder.set_decoder(self.sae_decoder)
+
+        if hasattr(self.sae_encoder, '_init_decoder_weights'):
+            self.sae_encoder._init_decoder_weights(
+                self.sae_decoder.W_dec
+            )
+
+        self.tied = tied
+
+        # get the number of concepts from the output of the sae_encoder
+        self.num_concepts = sae_encoder.out_channels
+
+    def forward(self, x):
+        """
+        Forward pass through SAE. Expects already-encoded features [B, act_size].
+        Encoding (audio_encoder / text_encoder) is done at training/validation step level.
+        """
+        if self.tied:
+            if hasattr(self.sae_encoder, '_init_decoder_weights'):
+                self.sae_encoder._init_decoder_weights(
+                    self.sae_decoder.W_dec
+                )
+        pre, z = self.sae_encoder(x)    
+        xhat = self.sae_decoder(z)
+        return x, z, xhat, pre
+
+    @torch.no_grad()
+    def inference(self, x, matryoshka_k=-1):
+        """
+        Inference pass. Expects already-encoded features [B, act_size].
+        Caller must encode with audio_encoder or text_encoder before calling.
+        """
+        pre, z = self.sae_encoder(x)
+        xhat = self.sae_decoder(z)
+        return x, z, xhat, pre
+
+class LightningSAE(SAE, L.LightningModule):
+
+    def __init__(self,
+                audio_encoder,
+                text_encoder,
+                sae_encoder,
+                sae_decoder,
+                optimizer = None,
+                scheduler = None,
+                loss_fn = None,
+                ckpt_path = None,
+                freeze = False,
+                preextracted_features = False,
+                tied = False,
+                data_mixing_ratio = 0.5,
+                **kwargs):
+        L.LightningModule.__init__(self)
+        super(LightningSAE, self).__init__(
+            audio_encoder=audio_encoder,
+            text_encoder=text_encoder,
+            sae_encoder=sae_encoder,
+            sae_decoder=sae_decoder,
+            ckpt_path=ckpt_path,
+            freeze=freeze, 
+            tied=tied,
+            **kwargs
+        )
+        self.optimizer = optimizer
+        self.scheduler = scheduler
+        self.loss_fn = loss_fn
+        self.preextracted_features = preextracted_features
+        self.reset_activations_dicts()
+        self.data_mixing_ratio = data_mixing_ratio
+
+    def state_dict(self, destination=None, prefix='', keep_vars=False):
+        """
+        Override state_dict to only save SAE encoder/decoder weights,
+        excluding audio_encoder and text_encoder.
+        """
+        state_dict = super().state_dict(destination=destination, prefix=prefix, keep_vars=keep_vars)
+        
+        # Filter out audio_encoder and text_encoder parameters
+        filtered_state_dict = {}
+        for key, value in state_dict.items():
+            if not (key.startswith('audio_encoder.') or key.startswith('text_encoder.')):
+                filtered_state_dict[key] = value
+
+        if hasattr(self, 'bridges'):
+            filtered_state_dict['bridges'] = self.bridges
+        else:
+            logging.warning("No bridges found in state_dict")
+        
+        return filtered_state_dict
+
+    def load_state_dict(self, state_dict, strict=True):
+        """
+        Override load_state_dict to only load SAE encoder/decoder weights.
+        """
+        # Restore non-parameter bridge artifacts saved by callbacks.
+        # These are plain dicts, not nn.Parameters/buffers, so we need to set them manually.
+        bridges = state_dict.pop('bridges', None)
+        if bridges is None:
+            logging.warning("No bridges found in state_dict")
+        if bridges is not None:
+            self.bridges = bridges
+
+        # Filter out audio_encoder and text_encoder from the state_dict if present
+        filtered_state_dict = {}
+        for key, value in state_dict.items():
+            if not (key.startswith('audio_encoder.') or key.startswith('text_encoder.')):
+                filtered_state_dict[key] = value
+        
+        return super().load_state_dict(filtered_state_dict, strict=strict)
+
+    def reset_activations_dicts(self):
+        """Reset activation and embedding dictionaries for all dataloaders."""
+        # Initialize empty dicts; will be populated per dataloader_idx
+        self.val_activations = {}
+        self.test_activations = {}
+        self.val_embeddings = {}
+        self.test_embeddings = {}
+    
+    def _get_dataset_name(self, dataloader_idx: int, mode: str) -> str:
+        """Get dataset name for a dataloader from datamodule.names (same logic as callbacks)."""
+        dm = getattr(self.trainer, "datamodule", None)
+        if dm is not None and hasattr(dm, "names"):
+            names = dm.names
+            if isinstance(names, dict) and mode in names:
+                name = names[mode].get(dataloader_idx)
+                if name is not None:
+                    return str(name)
+        return f"dataloader_{dataloader_idx}"
+
+    def _get_or_init_dataloader_storage(self, dataloader_idx, mode='val'):
+        """Get or initialize the activation and embedding dicts for a specific dataloader_idx."""
+        act_store = self.val_activations if mode == 'val' else self.test_activations
+        emb_store = self.val_embeddings if mode == 'val' else self.test_embeddings
+        if dataloader_idx not in act_store:
+            act_store[dataloader_idx] = {'audio': [], 'text': []}
+        if dataloader_idx not in emb_store:
+            emb_store[dataloader_idx] = {'audio': [], 'text': []}
+        return act_store[dataloader_idx], emb_store[dataloader_idx]
+
+    def _forward_and_loss(self, encoded, loss_kw):
+        """Forward pass + loss computation. Returns (loss_dict, activations, pre_penalty_activations)."""
+        _, activations, reconstruction, pre_penalty_activations = self(encoded)
+        loss_kw = {**loss_kw, "pre_penalty_activations": pre_penalty_activations}
+        loss_dict = self.loss_fn(reconstruction, activations, encoded, **loss_kw)
+        return loss_dict, activations, pre_penalty_activations
+
+    def _log_loss_dict(self, loss_dict, stage, modality=None):
+        """Log all keys in loss_dict with appropriate prefixes."""
+        for key, val in loss_dict.items():
+            p = LOSS_KEY_TO_PREFIX.get(key, key)
+            suffix = f"/{modality}" if modality else ""
+            self.log(f'{p}/{stage}{suffix}', val, prog_bar=True, on_step=True, sync_dist=True)
+
+    def training_step(self, batch, batch_idx):
+        audio, text = batch['audio'], batch['prompt']
+        has_audio = audio is not None and self.data_mixing_ratio <1.0
+        has_text = text is not None and self.data_mixing_ratio >0.0
+        loss_kw = dict(
+            num_batches_not_active=getattr(self.sae_encoder, 'num_batches_not_active', None),
+            decoder=self.sae_decoder,
+        )
+
+        # Encode inputs
+        encoded_audio = (self.audio_encoder(audio) if not self.preextracted_features else audio) if has_audio else None
+        encoded_text = self.text_encoder(text) if has_text else None
+
+        if has_audio and has_text:
+            # Half-and-half: take half of each modality
+            ratio = self.data_mixing_ratio
+            n_text = int(len(text) * ratio) # if ratio = 1, only text is used
+            n_audio = len(audio) - n_text # if ratio = 0, only audio is used
+            encoded_audio, encoded_text = encoded_audio[:n_audio], encoded_text[:n_text]
+            encoded = torch.cat([encoded_audio, encoded_text], dim=0)
+            n_audio = encoded_audio.size(0)
+
+            _, activations, reconstruction, pre_penalty_activations = self(encoded)
+
+            # Split outputs and compute losses per modality
+            def split(t): return (t[:n_audio], t[n_audio:]) if t is not None else (None, None)
+            enc_a, enc_t = split(encoded)
+            act_a, act_t = split(activations)
+            rec_a, rec_t = split(reconstruction)
+            pre_a, pre_t = split(pre_penalty_activations)
+            modality_losses = {}
+            if enc_a is not None and enc_a.size(0) > 0:
+                modality_losses["audio"] = self.loss_fn(
+                    rec_a, act_a, enc_a, **{**loss_kw, "pre_penalty_activations": pre_a}
+                )
+                self._log_loss_dict(modality_losses["audio"], "Train", "audio")
+            if enc_t is not None and enc_t.size(0) > 0:
+                modality_losses["text"] = self.loss_fn(
+                    rec_t, act_t, enc_t, **{**loss_kw, "pre_penalty_activations": pre_t}
+                )
+                self._log_loss_dict(modality_losses["text"], "Train", "text")
+
+            if not modality_losses:
+                raise ValueError(
+                    "Both audio and text slices are empty after applying data_mixing_ratio "
+                    f"({self.data_mixing_ratio})."
+                )
+
+            # Log averaged total over modalities that are actually present.
+            first_loss_dict = next(iter(modality_losses.values()))
+            for key in first_loss_dict:
+                vals = [loss_dict[key] for loss_dict in modality_losses.values() if key in loss_dict]
+                if not vals:
+                    continue
+                p = LOSS_KEY_TO_PREFIX.get(key, key)
+                self.log(
+                    f'{p}/Train/total',
+                    sum(vals) / len(vals),
+                    prog_bar=True,
+                    on_step=True,
+                    sync_dist=True,
+                )
+            total_loss = sum(loss_dict["loss"] for loss_dict in modality_losses.values()) / len(modality_losses)
+        else:
+            # Single modality
+            encoded = encoded_audio if has_audio else encoded_text
+            modality = "audio" if has_audio else "text"
+            loss_dict, _, _ = self._forward_and_loss(encoded, loss_kw)
+            self._log_loss_dict(loss_dict, "Train", modality)
+            total_loss = loss_dict['loss']
+
+        self.log('loss/Train/total', total_loss, prog_bar=True, on_step=True, sync_dist=True)
+        return total_loss
+    
+    def _eval_step(self, batch, batch_idx, dataloader_idx, stage):
+        """Shared logic for validation_step and test_step."""
+        audio, text = batch['audio'], batch['prompt']
+        dataset_name = self._get_dataset_name(dataloader_idx, stage)
+        acts, embs = self._get_or_init_dataloader_storage(dataloader_idx, mode=stage)
+        loss_kw = dict(
+            num_batches_not_active=getattr(self.sae_encoder, 'num_batches_not_active', None),
+            decoder=self.sae_decoder,
+        )
+        loss_dicts = {}
+
+        # Process each modality
+        for modality, data, encoder in [
+            ('audio', audio, lambda x: self.audio_encoder(x) if not self.preextracted_features else x),
+            ('text', text, self.text_encoder),
+        ]:
+            if data is None:
+                continue
+            encoded = encoder(data)
+            loss_dict, activations, _ = self._forward_and_loss(encoded, loss_kw)
+            loss_dicts[modality] = loss_dict
+            acts[modality].append(activations)
+            embs[modality].append(encoded.detach())
+            for key, val in loss_dict.items():
+                p = LOSS_KEY_TO_PREFIX.get(key, key)
+                self.log(f'{p}/{stage}/{dataset_name}/{modality}', val, sync_dist=True)
+
+        # Compute total loss
+        if not loss_dicts:
+            raise ValueError(f"Both audio and text are None in {stage} batch")
+        if len(loss_dicts) == 2:
+            for key in loss_dicts['audio']:
+                p = LOSS_KEY_TO_PREFIX.get(key, key)
+                avg = (loss_dicts['audio'][key] + loss_dicts['text'][key]) / 2
+                self.log(f'{p}/{stage}/{dataset_name}', avg, sync_dist=True)
+            total_loss = (loss_dicts['audio']['loss'] + loss_dicts['text']['loss']) / 2
+        else:
+            total_loss = next(iter(loss_dicts.values()))['loss']
+
+        self.log(f'loss/{stage}/{dataset_name}', total_loss, sync_dist=True)
+        return total_loss
+
+    def validation_step(self, batch, batch_idx, dataloader_idx=0):
+        return self._eval_step(batch, batch_idx, dataloader_idx, 'val')
+
+    def test_step(self, batch, batch_idx, dataloader_idx=0):
+        return self._eval_step(batch, batch_idx, dataloader_idx, 'test')
+
+    def on_validation_epoch_start(self):
+        super().on_validation_epoch_start()
+        self.reset_activations_dicts()
+    
+    def on_test_epoch_start(self):
+        super().on_test_epoch_start()
+        self.reset_activations_dicts()
+
+    # def on_before_optimizer_step(self, optimizer):
+        # Ensure decoder directions stay unit-norm (TopK paper recipe)
+        # if hasattr(self.sae_decoder, "make_decoder_weights_and_grad_unit_norm"):
+        #     self.sae_decoder.make_decoder_weights_and_grad_unit_norm()
+
+    def on_before_optimizer_step(self, optimizer):
+        with torch.no_grad():
+            self.sae_decoder.W_dec.data = self.sae_decoder.W_dec.data / (self.sae_decoder.W_dec.data.norm(dim=-1, keepdim=True) + 1e-8)
+
+
+    

steerable_retrieval/models/utils/__init__.py

@@ -0,0 +1,19 @@
+"""Model utilities."""
+
+from steerable_retrieval.models.utils.losses import (
+    SAELoss,
+    VanillaSAELoss,
+    TopKSAELoss,
+    BatchTopKSAELoss,
+    JumpReLUSAELoss,
+    MatryoshkaBatchTopKSAELoss,
+)
+
+__all__ = [
+    'SAELoss',
+    'VanillaSAELoss',
+    'TopKSAELoss',
+    'BatchTopKSAELoss',
+    'JumpReLUSAELoss',
+    'MatryoshkaBatchTopKSAELoss',
+]

steerable_retrieval/models/utils/losses.py

@@ -0,0 +1,398 @@
+"""Loss functions for SAE training."""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from steerable_retrieval.models.sae.penalties import StepFunction
+
+
+def r2_score(reconstruction: torch.Tensor, target: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
+    """Coefficient of determination (variance-weighted, aggregated over all outputs).
+
+    ``1 - SS_res / SS_tot`` with ``target`` as ground truth. Logged as a reconstruction
+    quality metric during SAE training.
+    """
+    reconstruction = reconstruction.float()
+    target = target.float()
+    ss_res = (target - reconstruction).pow(2).sum()
+    ss_tot = (target - target.mean(dim=0, keepdim=True)).pow(2).sum()
+    return 1.0 - ss_res / (ss_tot + eps)
+
+
+class SAELoss(nn.Module):
+    """
+    Base loss function for SAE training.
+    Computes loss from reconstruction, activations, and target.
+    """
+
+    def __init__(self, l1_coeff=1e-3, l2_coeff=1.0):
+        super().__init__()
+        self.l1_coeff = l1_coeff
+        self.l2_coeff = l2_coeff
+
+    def forward(self, reconstruction, activations, target, matryoshka=None, num_batches_not_active=None, decoder=None, **kwargs):
+        """
+        Compute SAE loss.
+
+        Args:
+            reconstruction: Reconstructed features [batch, ..., act_size]
+            activations: Feature activations [batch, ..., dict_size]
+            target: Target features [batch, ..., act_size]
+            matryoshka: Matryoshka module (optional, for compatibility)
+            num_batches_not_active: Tracking of inactive features (optional, for compatibility)
+            decoder: SAE decoder (optional, used by TopK/BatchTopK aux loss)
+
+        Returns:
+            Dictionary with loss components
+        """
+        # L2 reconstruction loss (MSE scaled by l2_coeff)
+        rec_err = (reconstruction.float() - target.float()).pow(2)
+        mse = rec_err.mean()
+        l2_loss = self.l2_coeff * mse
+
+        r2 = r2_score(reconstruction, target)
+        # dead_codes_frac = dead_codes(activations)
+        # ndn = dead_codes_frac.mean() if hasattr(dead_codes_frac, "mean") else dead_codes_frac
+
+        alive = (activations > 0).any(dim=0)
+        dead_ratio = 1 - alive.float().mean()
+
+        # L1 sparsity loss
+        l1_norm = activations.float().abs().sum(-1).mean()
+        l1_loss = self.l1_coeff * l1_norm
+
+        # L0 norm (number of active features)
+        l0_norm = (activations > 0).float().sum(-1).mean()
+
+        # Total loss
+        loss = l2_loss + l1_loss
+
+        return {
+            'loss': loss,
+            'l2_loss': l2_loss,
+            'l1_loss': l1_loss,
+            'l0_norm': l0_norm,
+            'l1_norm': l1_norm,
+            'mse': mse.detach(),
+            'r2': r2,
+            'ndn': dead_ratio.detach(),
+        }
+
+
+class VanillaSAELoss(SAELoss):
+    """Loss for VanillaSAE - standard L1 + L2."""
+    pass
+
+
+class TopKSAELoss(SAELoss):
+    """Loss for TopKSAE - includes auxiliary loss for dead features."""
+
+    def __init__(self, l1_coeff=1e-3, l2_coeff=1.0, aux_penalty=1.0, top_k_aux=10, n_batches_to_dead=10, orthogonal_loss=0.0):
+        super().__init__(l1_coeff, l2_coeff)
+        self.aux_penalty = aux_penalty
+        self.top_k_aux = top_k_aux
+        self.n_batches_to_dead = n_batches_to_dead
+        self.orthogonal_coeff = orthogonal_loss
+    def forward(self, reconstruction, activations, target,
+            num_batches_not_active=None, matryoshka=None, decoder=None,
+            pre_penalty_activations=None):
+        loss_dict = super().forward(reconstruction, activations, target)
+        aux_loss = self._get_auxiliary_loss(
+            target=target,
+            reconstruction=reconstruction,
+            activations=activations,
+            # num_batches_not_active=num_batches_not_active,
+            decoder=decoder,
+            pre_penalty_activations=pre_penalty_activations,
+        )
+        loss_dict['aux_loss'] = aux_loss
+        loss_dict['orthogonal_loss'] = self.offdiag_gram_loss(D = decoder.W_dec)
+        loss_dict['loss'] = loss_dict['loss'] + aux_loss + loss_dict['orthogonal_loss'] * self.orthogonal_coeff
+        return loss_dict
+    
+
+    def _get_auxiliary_loss(
+        self,
+        target,
+        reconstruction,
+        activations,
+        decoder=None,
+        pre_penalty_activations=None,
+    ):
+        # Match their behavior: if missing inputs, aux=0
+        if decoder is None or pre_penalty_activations is None or self.aux_penalty <= 0:
+            return torch.tensor(0.0, device=target.device, dtype=target.dtype)
+
+        # IMPORTANT: order matters (their comment) => residual = x - x_hat
+        residual = (target - reconstruction)
+
+        # Their AuxK uses relu(pre) and removes chosen codes by subtracting post-TopK codes
+        aux_src = F.relu(pre_penalty_activations) - activations
+
+        # Choose top half of *non-chosen* activations (after subtraction)
+        k = max(1, aux_src.shape[-1] // 2)
+        topk = torch.topk(aux_src, k=k, dim=-1)
+
+        aux_codes = torch.zeros_like(aux_src).scatter(-1, topk.indices, topk.values)
+
+        # Predict residual using full dictionary; no bias term
+        residual_hat = aux_codes @ decoder.W_dec  # [B, act_size]
+
+        aux_mse = (residual - residual_hat).pow(2).mean()
+        return self.aux_penalty * aux_mse
+
+
+    @staticmethod
+    def offdiag_gram_loss(
+        D: torch.Tensor,
+        atoms_dim: int | None = None,
+        normalize: bool = True,
+        squared: bool = True,
+        reduction: str = "mean",
+        eps: float = 1e-8,
+    ) -> torch.Tensor:
+        """
+        Penalize off-diagonal entries of the (optionally normalized) Gram matrix.
+
+        This encourages *incoherence* (low pairwise cosine similarity) between atoms,
+        without requiring strict orthogonality (which is infeasible when overcomplete).
+
+        Parameters
+        ----------
+        D : Tensor
+            Dictionary/decoder matrix, shape [A, d] or [d, A] where A=#atoms/features.
+            Examples:
+            - If W_dec is [dict_size, act_size], then atoms_dim=0.
+            - If W_dec is [act_size, dict_size], then atoms_dim=1.
+        atoms_dim : {0,1} or None
+            Which dimension indexes atoms. If None, we infer it by assuming atoms are
+            the larger dimension (common in overcomplete SAEs).
+        normalize : bool
+            If True, L2-normalize each atom before computing Gram (cosine coherence).
+        squared : bool
+            If True use L2 penalty on off-diagonals (sum of squares). If False use L1.
+        reduction : {"mean","sum"}
+            How to reduce over off-diagonal entries.
+        eps : float
+            Numerical stability for normalization.
+
+        Returns
+        -------
+        loss : Tensor
+            Scalar tensor.
+        """
+        A = D.clone()
+
+        if normalize:
+            A = A / (A.norm(dim=1, keepdim=True).clamp_min(eps))
+
+        # Gram: [A, A]
+        G = A @ A.t()
+
+        # Off-diagonal mask
+        n = G.shape[0]
+        if n <= 1:
+            return G.new_zeros(())
+        off = ~torch.eye(n, dtype=torch.bool, device=G.device)
+
+        vals = G[off]
+        if squared:
+            vals = vals * vals
+        else:
+            vals = vals.abs()
+
+        if reduction == "mean":
+            return vals.mean()
+        elif reduction == "sum":
+            return vals.sum()
+        else:
+            raise ValueError("reduction must be 'mean' or 'sum'")
+
+
+class BatchTopKSAELoss(TopKSAELoss):
+    """Loss for BatchTopKSAE - same as TopK but with batch-level sparsity."""
+    pass
+
+
+class JumpReLUSAELoss(SAELoss):
+    """Loss for JumpReLUSAE - uses L0 norm instead of L1."""
+    
+    def __init__(self, l1_coeff=1e-3, l2_coeff=1.0, bandwidth=0.1):
+        super().__init__(l1_coeff, l2_coeff)
+        self.bandwidth = bandwidth
+    
+    def forward(self, reconstruction, activations, target, log_threshold=None, matryoshka=None, num_batches_not_active=None, decoder=None):
+        """
+        Compute JumpReLU SAE loss.
+        
+        Args:
+            reconstruction: Reconstructed features
+            activations: Feature activations
+            target: Target features
+            log_threshold: Log threshold parameter from JumpReLU (optional)
+            matryoshka: Matryoshka module (optional, for compatibility)
+            num_batches_not_active: Tracking of inactive features (optional, for compatibility)
+        """
+        # L2 reconstruction loss
+        rec_err = (reconstruction.float() - target.float()).pow(2)
+        mse = rec_err.mean()
+        l2_loss = self.l2_coeff * mse
+        target_centered = target.float() - target.float().mean(dim=0)
+        baseline_mse = (target_centered.pow(2)).mean().clamp_min(1e-12)
+        fvu = (mse / baseline_mse).detach()
+        ndn = (activations.abs().sum(dim=0) == 0).float().mean().detach()
+
+        # L0 norm using StepFunction if log_threshold is provided
+        if log_threshold is not None:
+            l0 = StepFunction.apply(activations, log_threshold, self.bandwidth).sum(dim=-1).mean()
+        else:
+            l0 = (activations > 0).float().sum(-1).mean()
+
+        l0_loss = self.l1_coeff * l0
+        l1_loss = l0_loss  # For JumpReLU, L1 loss is the same as L0 loss
+
+        loss = l2_loss + l1_loss
+
+        return {
+            'loss': loss,
+            'l2_loss': l2_loss,
+            'l1_loss': l1_loss,
+            'l0_norm': l0,
+            'l1_norm': l0,
+            'mse': mse.detach(),
+            'fvu': fvu,
+            'ndn': ndn,
+        }
+
+class MatryoshkaBatchTopKSAELoss(SAELoss):
+    """Loss for Matryoshka BatchTopK SAE with progressive reconstruction losses.
+
+    Computes intermediate reconstructions by progressively decoding each
+    feature group through ``decoder.W_dec`` slices.  Supports two weighting
+    strategies from Zaigrajew et al. (2025):
+
+    * **uniform** (UW): ``alpha_i = 1`` for all granularity levels.
+    * **reverse** (RW): ``alpha_i = h - i + 1``, giving higher weight to
+      sparser (earlier) levels so the model prioritises reconstruction
+      quality at low feature counts.
+
+    Also includes a TopK-style auxiliary loss for dead features.
+    """
+
+    def __init__(self, group_sizes, l1_coeff=1e-3, l2_coeff=1.0,
+                 aux_penalty=1.0, top_k_aux=10, n_batches_to_dead=10,
+                 weighting='uniform'):
+        super().__init__(l1_coeff, l2_coeff)
+        self.group_sizes = group_sizes
+        self.group_indices = [0] + torch.cumsum(torch.tensor(group_sizes), dim=0).tolist()
+        self.aux_penalty = aux_penalty
+        self.top_k_aux = top_k_aux
+        self.n_batches_to_dead = n_batches_to_dead
+
+        h = len(group_sizes)
+        if weighting == 'uniform':
+            self.alphas = [1.0] * h
+        elif weighting == 'reverse':
+            self.alphas = [float(h - i) for i in range(h)]
+        else:
+            raise ValueError(f"Unknown weighting strategy '{weighting}'. Use 'uniform' or 'reverse'.")
+        self.alpha_sum = sum(self.alphas)
+
+    def forward(self, reconstruction, activations, target,
+                decoder=None, pre_penalty_activations=None,
+                num_batches_not_active=None, **kwargs):
+        """
+        Compute Matryoshka BatchTopK loss with progressive intermediate losses.
+
+        Args:
+            reconstruction: Final reconstructed features [batch, act_size]
+            activations: Sparse feature activations [batch, dict_size]
+            target: Target features [batch, act_size]
+            decoder: SAE decoder (required -- provides W_dec and b_dec)
+            pre_penalty_activations: Pre-activation features before ReLU/TopK [batch, dict_size]
+            num_batches_not_active: Per-feature inactivity counter (optional)
+
+        Returns:
+            Dictionary with loss components.
+        """
+        target_f = target.float()
+
+        # -- Intermediate reconstruction losses (progressive decoding) -------
+        b_dec = decoder.b_dec if decoder is not None else torch.zeros_like(target_f[0])
+
+        # Baseline: reconstruction using only b_dec (no features active)
+        baseline_l2 = self.l2_coeff * (b_dec - target_f).pow(2).mean()
+
+        x_recon = b_dec  # accumulates progressive reconstruction
+        intermediate_l2s = []
+        for i in range(len(self.group_sizes)):
+            start = self.group_indices[i]
+            end = self.group_indices[i + 1]
+            x_recon = activations[..., start:end] @ decoder.W_dec[start:end] + x_recon
+            intermediate_l2s.append(
+                self.l2_coeff * (x_recon.float() - target_f).pow(2).mean()
+            )
+
+        # Weighted mean L2: baseline + weighted intermediates, divided by
+        # (1 + sum_of_weights) to match the reference's (h+1) averaging.
+        weighted_l2 = baseline_l2 + sum(
+            a * l for a, l in zip(self.alphas, intermediate_l2s)
+        )
+        mean_l2 = weighted_l2 / (1.0 + self.alpha_sum)
+
+        l2_stack = torch.stack(intermediate_l2s)
+        min_l2 = l2_stack.min()
+        max_l2 = l2_stack.max()
+
+        # -- Final-reconstruction metrics (mse, r2, ndn) --------------------
+        mse = (reconstruction.float() - target_f).pow(2).mean()
+        r2 = r2_score(reconstruction, target)
+        alive = (activations > 0).any(dim=0)
+        ndn = 1 - alive.float().mean()
+
+        # -- Sparsity losses -------------------------------------------------
+        l1_norm = activations.float().abs().sum(-1).mean()
+        l1_loss = self.l1_coeff * l1_norm
+        l0_norm = (activations > 0).float().sum(-1).mean()
+
+        # -- Auxiliary loss for dead features ---------------------------------
+        aux_loss = self._get_auxiliary_loss(
+            target=target,
+            reconstruction=reconstruction,
+            activations=activations,
+            decoder=decoder,
+            pre_penalty_activations=pre_penalty_activations,
+        )
+
+        loss = mean_l2 + l1_loss + aux_loss
+
+        return {
+            'loss': loss,
+            'l2_loss': mean_l2,
+            'min_l2_loss': min_l2.detach(),
+            'max_l2_loss': max_l2.detach(),
+            'l1_loss': l1_loss,
+            'l0_norm': l0_norm,
+            'l1_norm': l1_norm,
+            'aux_loss': aux_loss,
+            'mse': mse.detach(),
+            'r2': r2,
+            'ndn': ndn.detach(),
+        }
+
+    def _get_auxiliary_loss(self, target, reconstruction, activations,
+                            decoder=None, pre_penalty_activations=None):
+        """TopK-style auxiliary loss: predict the residual using dead features."""
+        if decoder is None or pre_penalty_activations is None or self.aux_penalty <= 0:
+            return torch.tensor(0.0, device=target.device, dtype=target.dtype)
+
+        residual = target - reconstruction
+        aux_src = F.relu(pre_penalty_activations) - activations
+
+        k = max(1, aux_src.shape[-1] // 2)
+        topk = torch.topk(aux_src, k=k, dim=-1)
+        aux_codes = torch.zeros_like(aux_src).scatter(-1, topk.indices, topk.values)
+
+        residual_hat = aux_codes @ decoder.W_dec
+        aux_mse = (residual - residual_hat).pow(2).mean()
+        return self.aux_penalty * aux_mse

steerable_retrieval/models/utils/schedulers.py

@@ -0,0 +1,44 @@
+#write a pytorch learning rate scheduler for cosine decay with linear warmup
+
+from torch.optim import Optimizer
+import torch
+import math
+import torch.optim as optim
+import torch.nn as nn
+from tqdm import tqdm
+
+class CosineDecayWithLinearWarmup(torch.optim.lr_scheduler._LRScheduler):
+    def __init__(self, optimizer, warmup_steps = 5000, max_steps = 200000, base_lr = 0.0001, final_lr = 1e-6, last_step=-1):
+        self.warmup_steps = warmup_steps
+        self.max_steps = max_steps
+        self.base_lr = base_lr
+        self.final_lr = final_lr
+        self.last_step = last_step
+        super(CosineDecayWithLinearWarmup, self).__init__(optimizer, last_step)
+        
+    def get_lr(self):
+        if self.last_step < self.warmup_steps:
+            #warmup from 0 to base_lr
+            return [self.last_step / self.warmup_steps * base_lr for base_lr in self.base_lrs]
+        else:
+            return [self.final_lr + 0.5 * (base_lr - self.final_lr) * (1 + math.cos(math.pi * (self.last_step - self.warmup_steps) / (self.max_steps - self.warmup_steps))) for base_lr in self.base_lrs]
+    
+    def step(self, step=None):
+        if step is None:
+            step = self.last_step + 1
+        self.last_step = step
+        for param_group, lr in zip(self.optimizer.param_groups, self.get_lr()):
+            param_group['lr'] = lr
+
+# write a dummy test for the scheduler with 5k warmup, 200k max steps, a dummy network and optimizer, and return the step and lr history
+    @staticmethod
+    def test():
+        net = nn.Linear(10, 1)
+        optimizer = optim.Adam(net.parameters(), lr=0.0001)
+        scheduler = CosineDecayWithLinearWarmup(optimizer, warmup_steps = 5000, max_steps = 200000, base_lr = 0.0001, final_lr = 1e-6)
+        step_lr_history = []
+        for step in tqdm(range(0, 200000)):
+            optimizer.step()
+            scheduler.step()
+            step_lr_history.append((step, optimizer.param_groups[0]['lr']))
+        return step_lr_history

steerable_retrieval/steer/__init__.py

@@ -0,0 +1,31 @@
+"""Sparse steerable retrieval -- public API.
+
+Open-vocabulary concept control for dense music retrieval via sparse inversion in a
+trained SAE. See :class:`~steerable_retrieval.steer.slider.Slider` for the main entry point.
+"""
+
+from steerable_retrieval.steer.inversion import (
+    InversionResult,
+    MahalanobisPrior,
+    fit_mahalanobis_prior,
+    invert_concept,
+    load_default_prior,
+    mahalanobis_distance,
+)
+from steerable_retrieval.steer.retrieval import topk_cosine_neighbors
+from steerable_retrieval.steer.steering import build_edit_mask, idf_weights, steer_sparse
+from steerable_retrieval.steer.slider import Slider
+
+__all__ = [
+    "Slider",
+    "InversionResult",
+    "MahalanobisPrior",
+    "fit_mahalanobis_prior",
+    "load_default_prior",
+    "invert_concept",
+    "mahalanobis_distance",
+    "build_edit_mask",
+    "idf_weights",
+    "steer_sparse",
+    "topk_cosine_neighbors",
+]

steerable_retrieval/steer/inversion.py

@@ -0,0 +1,389 @@
+"""Sparse inversion for open-vocabulary concept attribution.
+
+This is the core method of *"Steering dense music retrieval with open-vocabulary
+concept discovery"* (ISMIR 2026). Given a free-form concept's text embedding
+``z_c`` in a joint music--text space, we recover a **sparse code** whose decoded
+audio embedding reconstructs ``z_c`` while staying close to the empirical audio
+manifold (a Mahalanobis prior). Two solvers are provided:
+
+* ``method="adam"``  -- gradient descent on the latent pre-activations through the
+  SAE's own sparsifying operator (the general, differentiable variant).
+* ``method="fista"`` -- a fast linear-inverse solver exploiting the linear decoder
+  (~20 ms/inversion on CPU; used for the interactive demo).
+
+The recovered support is the concept "slider": the set of sparse features that,
+when amplified or suppressed, steer retrieval along the concept axis.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+# --------------------------------------------------------------------------- #
+# Mahalanobis audio-manifold prior
+# --------------------------------------------------------------------------- #
+@dataclass
+class MahalanobisPrior:
+    """Gaussian prior over L2-normalized audio embeddings.
+
+    ``mean`` is ``[1, d]`` and ``precision`` is ``[d, d]`` (inverse covariance).
+    Used to keep inverted concepts near the region occupied by real audio.
+    """
+
+    mean: torch.Tensor
+    precision: torch.Tensor
+
+    def to(self, device=None, dtype=None) -> "MahalanobisPrior":
+        return MahalanobisPrior(
+            mean=self.mean.to(device=device, dtype=dtype),
+            precision=self.precision.to(device=device, dtype=dtype),
+        )
+
+    def save(self, path, *, meta: Optional[dict] = None) -> None:
+        """Save the prior as a compact ``.npz`` (mean, precision, optional JSON meta)."""
+        arrs = {
+            "mean": self.mean.detach().cpu().numpy().astype(np.float32),
+            "precision": self.precision.detach().cpu().numpy().astype(np.float32),
+        }
+        if meta:
+            import json
+
+            arrs["meta"] = np.frombuffer(json.dumps(meta).encode("utf-8"), dtype=np.uint8)
+        np.savez_compressed(path, **arrs)
+
+    @classmethod
+    def load(cls, path, *, device=None, dtype: torch.dtype = torch.float32) -> "MahalanobisPrior":
+        """Load a prior saved by :meth:`save`."""
+        with np.load(path, allow_pickle=False) as data:
+            mean = torch.from_numpy(data["mean"]).to(device=device, dtype=dtype)
+            precision = torch.from_numpy(data["precision"]).to(device=device, dtype=dtype)
+        if mean.dim() == 1:
+            mean = mean.unsqueeze(0)
+        return cls(mean=mean, precision=precision)
+
+
+def load_default_prior(name: str = "muq_mulan_music4all", *, device=None) -> MahalanobisPrior:
+    """Load a Mahalanobis prior shipped with the package (a library constant).
+
+    The default (``muq_mulan_music4all``) is fit once on Music4All MuQ-MuLan audio
+    embeddings and packaged under ``steerable_retrieval/assets/``. It is a fixed
+    distributional constant of the audio manifold — independent of whatever corpus
+    a caller later retrieves over.
+    """
+    from importlib import resources
+
+    res = resources.files("steerable_retrieval.assets").joinpath(f"{name}_prior.npz")
+    with resources.as_file(res) as p:
+        return MahalanobisPrior.load(p, device=device)
+
+
+def fit_mahalanobis_prior(
+    embeddings,
+    *,
+    cov_eps: float = 1e-3,
+    shrinkage: float = 0.05,
+    device=None,
+    dtype: torch.dtype = torch.float32,
+) -> MahalanobisPrior:
+    """Fit a shrinkage Mahalanobis prior to a set of audio embeddings.
+
+    Args:
+        embeddings: ``[N, d]`` array/tensor of audio embeddings (need not be normalized).
+        cov_eps: diagonal jitter added to the covariance for invertibility.
+        shrinkage: convex shrinkage toward a scaled identity (Ledoit-Wolf style).
+    """
+    x = np.asarray(
+        embeddings.detach().cpu().numpy() if isinstance(embeddings, torch.Tensor) else embeddings,
+        dtype=np.float64,
+    )
+    x = x / np.clip(np.linalg.norm(x, axis=1, keepdims=True), 1e-12, None)
+    mean = x.mean(axis=0, keepdims=True)
+    centered = x - mean
+    denom = max(int(centered.shape[0] - 1), 1)
+    cov = (centered.T @ centered) / float(denom)
+    trace_scale = float(np.trace(cov) / cov.shape[0]) if cov.shape[0] > 0 else 1.0
+    cov = (1.0 - shrinkage) * cov + shrinkage * trace_scale * np.eye(cov.shape[0])
+    cov = cov + cov_eps * np.eye(cov.shape[0])
+    precision = np.linalg.inv(cov)
+    return MahalanobisPrior(
+        mean=torch.from_numpy(mean).to(device=device, dtype=dtype),
+        precision=torch.from_numpy(precision).to(device=device, dtype=dtype),
+    )
+
+
+def mahalanobis_distance(x: torch.Tensor, prior: MahalanobisPrior) -> torch.Tensor:
+    """Squared Mahalanobis distance of (L2-normalized) ``x`` to the prior, per row."""
+    x = F.normalize(x, dim=-1)
+    delta = x - prior.mean
+    return torch.einsum("bi,ij,bj->b", delta, prior.precision, delta)
+
+
+# --------------------------------------------------------------------------- #
+# Re-applying the SAE's sparsifying operator Pi to arbitrary latent logits
+# --------------------------------------------------------------------------- #
+def apply_encoder_sparse_transform(
+    model, latent_logits: torch.Tensor, use_encoder_penalty: bool = True
+) -> torch.Tensor:
+    """Apply the SAE encoder's sparsifying nonlinearity ``Pi`` to ``latent_logits``.
+
+    This lets the Adam inversion optimize *pre-activations* while keeping the exact
+    training-time sparsity structure of the model (TopK / BatchTopK / JumpReLU / ReLU).
+    """
+    encoder = model.sae_encoder
+    encoder_name = type(encoder).__name__
+    if encoder_name == "VanillaSAEEncoder":
+        return F.relu(latent_logits)
+    if encoder_name == "TopKSAEEncoder":
+        acts = F.relu(latent_logits)
+        return encoder.penalty(acts) if use_encoder_penalty else acts
+    if encoder_name in {"BatchTopKSAEEncoder", "MatryoshkaBatchTopKSAEEncoder"}:
+        acts = F.relu(latent_logits)
+        if not use_encoder_penalty:
+            return acts
+        threshold = getattr(encoder, "running_threshold", None)
+        if threshold is None:
+            threshold = encoder.penalty._get_threshold(acts)
+        return encoder.penalty(acts, threshold)
+    if encoder_name == "JumpReLUSAEEncoder":
+        return encoder.jumprelu(latent_logits) if use_encoder_penalty else F.relu(latent_logits)
+    raise NotImplementedError(
+        f"Latent-logit inversion is not implemented for encoder type {encoder_name}."
+    )
+
+
+# --------------------------------------------------------------------------- #
+# FISTA solver for the linear inverse problem with a Mahalanobis term
+# --------------------------------------------------------------------------- #
+def soft_threshold(x: torch.Tensor, tau: float) -> torch.Tensor:
+    return torch.sign(x) * torch.clamp(torch.abs(x) - tau, min=0.0)
+
+
+def fista_mahalanobis(
+    A: torch.Tensor,
+    y: torch.Tensor,
+    mu: torch.Tensor,
+    M: torch.Tensor,
+    lam: float,
+    gamma: float,
+    n_iter: int = 200,
+    L: Optional[float] = None,
+    power_iters: int = 30,
+    enforce_nonnegative: bool = True,
+) -> torch.Tensor:
+    """Solve ``min_x 0.5||Ax - y||^2 + 0.5*gamma*(Ax-mu)^T M (Ax-mu) + lam||x||_1``.
+
+    ``A`` is ``W_dec^T`` (``[d, m]``); ``x`` is the sparse code (``[m]``). Non-negativity
+    matches ReLU-style SAE codes. ``L`` (Lipschitz constant) is estimated by power
+    iteration on the Hessian ``A^T A + gamma A^T M A`` when not provided.
+    """
+    n, m = A.shape
+    del n
+    x = torch.zeros(m, device=A.device, dtype=A.dtype)
+    z = x.clone()
+    t = 1.0
+    if L is None:
+        v = torch.randn(m, device=A.device, dtype=A.dtype)
+        v = v / v.norm().clamp_min(1e-12)
+        for _ in range(power_iters):
+            Hv = A.T @ (A @ v) + gamma * (A.T @ (M @ (A @ v)))
+            v = Hv / Hv.norm().clamp_min(1e-12)
+        Hv = A.T @ (A @ v) + gamma * (A.T @ (M @ (A @ v)))
+        L = torch.dot(v, Hv)
+    eta = 1.0 / float(torch.as_tensor(L).item())
+    for _ in range(n_iter):
+        Az = A @ z
+        grad = A.T @ (Az - y) + gamma * (A.T @ (M @ (Az - mu)))
+        x_new = soft_threshold(z - eta * grad, lam * eta)
+        if enforce_nonnegative:
+            x_new = F.relu(x_new)
+        t_new = (1 + (1 + 4 * t * t) ** 0.5) / 2
+        z = x_new + ((t - 1) / t_new) * (x_new - x)
+        x, t = x_new, t_new
+    return x
+
+
+# --------------------------------------------------------------------------- #
+# Inversion result + the two solvers + a dispatcher
+# --------------------------------------------------------------------------- #
+@dataclass
+class InversionResult:
+    """Output of a concept inversion.
+
+    ``sparse_values`` is the recovered sparse code over the SAE dictionary ``[m]``;
+    ``support_mask`` is the boolean support (active features). ``final_text_cosine``
+    and ``final_mahalanobis_distance`` are diagnostics of solution quality.
+    """
+
+    sparse_values: torch.Tensor
+    support_mask: torch.Tensor
+    final_text_cosine: float
+    final_mahalanobis_distance: float
+    method: str
+
+
+def _decoder_matrices(model, device, dtype):
+    """Return ``A = W_dec^T`` (``[d, m]``) and ``b_dec`` (``[d]``)."""
+    A = model.sae_decoder.W_dec.detach().to(device=device, dtype=dtype).T
+    b_dec = model.sae_decoder.b_dec.detach().to(device=device, dtype=dtype).view(-1)
+    return A, b_dec
+
+
+def invert_concept_adam(
+    model,
+    text_embedding: torch.Tensor,
+    prior: MahalanobisPrior,
+    *,
+    init_sparse: Optional[torch.Tensor] = None,
+    num_steps: int = 500,
+    lr: float = 1e-3,
+    reg_strength: float = 1e-4,
+    use_encoder_penalty: bool = True,
+    activation_threshold: float = 1e-6,
+) -> InversionResult:
+    """Adam inversion: optimize latent pre-activations through the SAE operator ``Pi``.
+
+    ``init_sparse`` (``[1, m]`` or ``[m]``) initializes the latent logits; the paper
+    initializes from the sparse code of the nearest audio neighbour to ``z_c`` for
+    stable convergence (see :func:`nearest_audio_sparse_init`).
+    """
+    device = text_embedding.device
+    dtype = text_embedding.dtype
+    prior = prior.to(device=device, dtype=dtype)
+    if init_sparse is None:
+        init = torch.zeros((1, model.sae_encoder.dict_size), device=device, dtype=dtype)
+    else:
+        init = init_sparse.view(1, -1).to(device=device, dtype=dtype)
+
+    latent_logits = torch.nn.Parameter(init.clone())
+    optimizer = torch.optim.AdamW([latent_logits], lr=lr)
+    for _ in range(num_steps):
+        optimizer.zero_grad(set_to_none=True)
+        z_sparse = apply_encoder_sparse_transform(model, latent_logits, use_encoder_penalty)
+        z_hat = model.sae_decoder(z_sparse)
+        text_cos = F.cosine_similarity(z_hat, text_embedding, dim=-1).mean()
+        mahal = mahalanobis_distance(z_hat, prior).mean()
+        loss = (1.0 - text_cos) + reg_strength * mahal
+        loss.backward()
+        optimizer.step()
+
+    with torch.no_grad():
+        z_sparse = apply_encoder_sparse_transform(model, latent_logits, use_encoder_penalty)
+        z_hat = model.sae_decoder(z_sparse)
+        sparse_values = z_sparse.squeeze(0).detach().cpu()
+        support_mask = (sparse_values.abs() > activation_threshold).bool()
+        final_cos = float(F.cosine_similarity(z_hat, text_embedding, dim=-1).mean().item())
+        final_mahal = float(mahalanobis_distance(z_hat, prior).mean().item())
+    return InversionResult(sparse_values, support_mask, final_cos, final_mahal, "adam")
+
+
+def invert_concept_fista(
+    model,
+    text_embedding: torch.Tensor,
+    prior: MahalanobisPrior,
+    *,
+    num_steps: int = 250,
+    l1_lambda: float = 0.01,
+    gamma: float = 1e-4,
+    lipschitz: Optional[float] = None,
+    power_iters: int = 30,
+    enforce_nonnegative: bool = True,
+    use_encoder_penalty: bool = True,
+    apply_encoder_penalty_after: bool = True,
+    activation_threshold: float = 1e-6,
+) -> InversionResult:
+    """FISTA inversion exploiting the linear decoder (fast; used by the live demo)."""
+    device = text_embedding.device
+    dtype = text_embedding.dtype
+    prior = prior.to(device=device, dtype=dtype)
+    A, b_dec = _decoder_matrices(model, device, dtype)
+    y = text_embedding.squeeze(0).detach().to(device=device, dtype=dtype) - b_dec
+    mu = prior.mean.squeeze(0).detach().to(device=device, dtype=dtype) - b_dec
+    M = prior.precision.detach().to(device=device, dtype=dtype)
+
+    x = fista_mahalanobis(
+        A, y, mu, M, l1_lambda, gamma,
+        n_iter=num_steps, L=lipschitz,
+        power_iters=power_iters, enforce_nonnegative=enforce_nonnegative,
+    )
+    z_sparse = x.unsqueeze(0)
+    if apply_encoder_penalty_after:
+        z_sparse = apply_encoder_sparse_transform(model, z_sparse, use_encoder_penalty)
+
+    with torch.no_grad():
+        z_hat = model.sae_decoder(z_sparse)
+        sparse_values = z_sparse.squeeze(0).detach().cpu()
+        support_mask = (sparse_values.abs() > activation_threshold).bool()
+        final_cos = float(F.cosine_similarity(z_hat, text_embedding, dim=-1).mean().item())
+        final_mahal = float(mahalanobis_distance(z_hat, prior).mean().item())
+    return InversionResult(sparse_values, support_mask, final_cos, final_mahal, "fista")
+
+
+def nearest_audio_sparse_init(
+    model,
+    text_embedding: torch.Tensor,
+    audio_embeddings: np.ndarray,
+    audio_norms: Optional[np.ndarray] = None,
+) -> torch.Tensor:
+    """Sparse code of the audio embedding nearest (cosine) to ``text_embedding``.
+
+    Recommended initialization for Adam inversion (paper Section 3.3).
+    """
+    text_norm = F.normalize(text_embedding.detach(), dim=-1).squeeze(0).cpu().numpy().astype(np.float32)
+    audio_embeddings = np.asarray(audio_embeddings, dtype=np.float32)
+    if audio_norms is None:
+        audio_norms = np.clip(np.linalg.norm(audio_embeddings, axis=1), 1e-12, None)
+    sims = (audio_embeddings @ text_norm) / audio_norms
+    best = int(np.argmax(sims))
+    device = model.sae_decoder.W_dec.device
+    batch = torch.from_numpy(audio_embeddings[best][None]).to(device)
+    with torch.inference_mode():
+        _, z, _, _ = model.inference(batch)
+    return z.detach().squeeze(0)
+
+
+def prior_mean_sparse_init(model, prior: MahalanobisPrior) -> torch.Tensor:
+    """Sparse code of the prior's mean embedding — a manifold-centered Adam init.
+
+    Used when no ``audio_embeddings`` are available for a nearest-neighbour init, so
+    Adam inversion still starts from a plausible audio point rather than zeros.
+    """
+    device = model.sae_decoder.W_dec.device
+    mean = prior.mean.to(device=device, dtype=torch.float32)
+    with torch.inference_mode():
+        _, z, _, _ = model.inference(mean)
+    return z.detach().squeeze(0)
+
+
+def invert_concept(
+    model,
+    text_embedding: torch.Tensor,
+    prior: MahalanobisPrior,
+    *,
+    method: str = "adam",
+    audio_embeddings: Optional[np.ndarray] = None,
+    **kwargs,
+) -> InversionResult:
+    """Dispatch to the requested inversion solver.
+
+    For ``method="adam"`` the optimization is initialized from the nearest audio
+    neighbour when ``audio_embeddings`` is provided (recommended by the paper), and
+    otherwise from the prior mean — so inversion needs only the prior, not a corpus.
+    """
+    method = method.lower()
+    if method == "adam":
+        if kwargs.get("init_sparse") is None:
+            kwargs["init_sparse"] = (
+                nearest_audio_sparse_init(model, text_embedding, audio_embeddings)
+                if audio_embeddings is not None
+                else prior_mean_sparse_init(model, prior)
+            )
+        return invert_concept_adam(model, text_embedding, prior, **kwargs)
+    if method == "fista":
+        return invert_concept_fista(model, text_embedding, prior, **kwargs)
+    raise ValueError(f"Unknown inversion method {method!r}; use 'adam' or 'fista'.")

steerable_retrieval/steer/loading.py

@@ -0,0 +1,131 @@
+"""Loading trained SAEs for the steering API.
+
+A :class:`SteerableSAE` is the minimal model surface the :class:`~steerable_retrieval.steer.slider.Slider`
+needs: a trained ``sae_encoder`` + ``sae_decoder`` and a ``text_encoder`` mapping concept
+strings to embeddings in the same joint space. This module resolves such a model from a
+checkpoint (local path or HuggingFace repo).
+
+NOTE: the released pretrained checkpoint (BatchTopK SAE on MuQ / music4all) is not published
+yet. Until then, construct a model in-memory and pass it via ``Slider(..., model=model)``.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+
+class SteerableSAE(nn.Module):
+    """Minimal SAE wrapper for inference-time steering (no training deps).
+
+    Args:
+        sae_encoder: a trained SAE encoder (``steerable_retrieval.models.sae.encoders``).
+        sae_decoder: a trained SAE decoder (``steerable_retrieval.models.sae.decoders``).
+        text_encoder: callable mapping ``list[str] -> [B, d]`` embeddings (e.g. the MuQ
+            text tower), living in the same joint space as the SAE's audio inputs.
+    """
+
+    def __init__(self, sae_encoder, sae_decoder, text_encoder):
+        super().__init__()
+        self.sae_encoder = sae_encoder
+        self.sae_decoder = sae_decoder
+        self.text_encoder = text_encoder
+        if not hasattr(self.sae_decoder, "b_dec"):
+            self.sae_decoder.b_dec = self.sae_encoder.b_dec
+
+    @torch.no_grad()
+    def inference(self, x):
+        """Encode dense features to a sparse code and back. Returns ``(x, z, xhat, pre)``."""
+        pre, z = self.sae_encoder(x)
+        xhat = self.sae_decoder(z)
+        return x, z, xhat, pre
+
+
+def _resolve_local_run(checkpoint: str, subfolder: Optional[str] = None):
+    """Return (ckpt_path, config_path) for a checkpoint, downloading from the Hub if
+    ``checkpoint`` is a ``org/repo`` id rather than a local path.
+
+    A Lightning run stores its Hydra config at ``<run>/.hydra/config.yaml`` and its
+    checkpoints under ``<run>/checkpoints/``. We use that config to rebuild the SAE
+    modules before loading the (encoder/decoder-only) weights. For a Hub repo carrying
+    several models, ``subfolder`` (e.g. ``"L0-20"``) selects which one.
+    """
+    import os
+
+    if os.path.exists(checkpoint):
+        ckpt_path = os.path.abspath(checkpoint)
+        run_dir = os.path.dirname(os.path.dirname(ckpt_path))  # .../checkpoints/x.ckpt -> run
+        cfg_path = os.path.join(run_dir, ".hydra", "config.yaml")
+        if not os.path.exists(cfg_path):
+            # allow a config.yaml sitting next to the checkpoint
+            alt = os.path.join(os.path.dirname(ckpt_path), "config.yaml")
+            cfg_path = alt if os.path.exists(alt) else cfg_path
+        return ckpt_path, cfg_path
+
+    # Otherwise treat it as a HuggingFace repo id: expects last.ckpt + config.yaml
+    # (optionally under `subfolder`).
+    from huggingface_hub import hf_hub_download
+
+    pre = f"{subfolder}/" if subfolder else ""
+    ckpt_path = hf_hub_download(checkpoint, filename=f"{pre}last.ckpt")
+    try:
+        cfg_path = hf_hub_download(checkpoint, filename=f"{pre}config.yaml")
+    except Exception:
+        cfg_path = hf_hub_download(checkpoint, filename=f"{pre}.hydra/config.yaml")
+    return ckpt_path, cfg_path
+
+
+def load_steerable_sae(
+    checkpoint: Optional[str],
+    *,
+    model_class: Optional[str] = None,
+    device: str = "cpu",
+    text_encoder=None,
+    build_text_encoder: bool = True,
+    config_path: Optional[str] = None,
+    subfolder: Optional[str] = None,
+) -> SteerableSAE:
+    """Resolve a :class:`SteerableSAE` from a trained Lightning checkpoint.
+
+    Args:
+        checkpoint: local path to a ``.ckpt`` (its run's ``.hydra/config.yaml`` is used
+            to rebuild the SAE modules), or a HuggingFace ``org/repo`` id carrying
+            ``last.ckpt`` + ``config.yaml``.
+        device: where to place the model.
+        text_encoder: a ready callable ``list[str] -> [B, d]``. If ``None`` and
+            ``build_text_encoder`` is True, the text tower from the run config (e.g.
+            MuQ-MuLan) is instantiated; if False, ``text_encoder`` stays ``None`` (useful
+            for steering/retrieval that never embeds new text).
+        config_path: override the run config path explicitly.
+    """
+    from hydra.utils import instantiate
+    from omegaconf import OmegaConf
+
+    ckpt_path, resolved_cfg = _resolve_local_run(checkpoint, subfolder=subfolder)
+    cfg = OmegaConf.load(config_path or resolved_cfg)
+
+    enc = instantiate(cfg.model.sae_encoder, device=device)
+    dec = instantiate(cfg.model.sae_decoder, device=device)
+    if not hasattr(dec, "b_dec"):
+        dec.b_dec = enc.b_dec
+
+    state = torch.load(ckpt_path, map_location=device)
+    sd = state.get("state_dict", state)
+    enc_sd = {k[len("sae_encoder."):]: v for k, v in sd.items() if k.startswith("sae_encoder.")}
+    dec_sd = {k[len("sae_decoder."):]: v for k, v in sd.items() if k.startswith("sae_decoder.")}
+    if not enc_sd or not dec_sd:
+        raise ValueError(
+            f"No sae_encoder/sae_decoder weights found in {ckpt_path}. "
+            f"Available prefixes: {sorted({k.split('.')[0] for k in sd})}"
+        )
+    enc.load_state_dict(enc_sd, strict=False)
+    dec.load_state_dict(dec_sd, strict=False)
+
+    if text_encoder is None and build_text_encoder:
+        text_encoder = instantiate(cfg.model.text_encoder, device=device)
+
+    model = SteerableSAE(enc, dec, text_encoder)
+    model.to(device).eval()
+    return model

steerable_retrieval/steer/prior_fit.py

@@ -0,0 +1,95 @@
+"""Fit and ship the default Mahalanobis manifold prior.
+
+The prior is a fixed distributional constant of the audio manifold, fit once on a
+sample of pre-extracted audio embeddings and packaged with the library so the
+:class:`~steerable_retrieval.steer.slider.Slider` works with no data from the caller.
+It is deliberately independent of any retrieval corpus.
+
+CLI (produces the packaged asset):
+
+    python -m steerable_retrieval.steer.prior_fit \
+        --manifest /path/to/embedding_lookup_manifest.csv \
+        --out steerable_retrieval/assets/muq_mulan_music4all_prior.npz \
+        --n-samples 20000
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import random
+from typing import Optional
+
+import numpy as np
+
+from steerable_retrieval.steer.inversion import MahalanobisPrior, fit_mahalanobis_prior
+
+
+def load_embeddings_from_manifest(
+    manifest_csv: str,
+    *,
+    n_samples: int = 20000,
+    split: Optional[str] = "train",
+    seed: int = 0,
+    path_col: str = "audio_embedding_path",
+    status_col: str = "audio_embedding_status",
+) -> np.ndarray:
+    """Sample ``n_samples`` audio-embedding ``.npy`` paths from a manifest and stack them."""
+    paths = []
+    with open(manifest_csv, newline="") as fh:
+        reader = csv.DictReader(fh)
+        for row in reader:
+            if split and row.get("split") != split:
+                continue
+            if status_col in row and row[status_col] not in ("ok", "", None):
+                continue
+            p = row.get(path_col)
+            if p:
+                paths.append(p)
+
+    rng = random.Random(seed)
+    if n_samples and len(paths) > n_samples:
+        paths = rng.sample(paths, n_samples)
+
+    embs = []
+    for p in paths:
+        try:
+            embs.append(np.asarray(np.load(p), dtype=np.float32).reshape(-1))
+        except Exception:
+            continue
+    if not embs:
+        raise RuntimeError(f"No embeddings loaded from {manifest_csv}")
+    return np.stack(embs, axis=0)
+
+
+def fit_and_save(
+    manifest_csv: str,
+    out_path: str,
+    *,
+    n_samples: int = 20000,
+    split: Optional[str] = "train",
+    model: str = "OpenMuQ/MuQ-MuLan-large",
+) -> MahalanobisPrior:
+    X = load_embeddings_from_manifest(manifest_csv, n_samples=n_samples, split=split)
+    prior = fit_mahalanobis_prior(X)
+    prior.save(
+        out_path,
+        meta={"model": model, "n_samples": int(X.shape[0]), "dim": int(X.shape[1]), "split": split, "source": "music4all"},
+    )
+    print(f"Fit prior on {X.shape[0]} embeddings (dim {X.shape[1]}) -> {out_path}")
+    return prior
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__)
+    ap.add_argument("--manifest", required=True, help="embedding manifest CSV")
+    ap.add_argument("--out", required=True, help="output .npz path")
+    ap.add_argument("--n-samples", type=int, default=20000)
+    ap.add_argument("--split", default="train")
+    ap.add_argument("--model", default="OpenMuQ/MuQ-MuLan-large")
+    args = ap.parse_args()
+    fit_and_save(args.manifest, args.out, n_samples=args.n_samples, split=args.split, model=args.model)
+
+
+if __name__ == "__main__":
+    main()

steerable_retrieval/steer/retrieval.py

@@ -0,0 +1,33 @@
+"""Cosine nearest-neighbour retrieval over a dense embedding corpus."""
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+
+
+def topk_cosine_neighbors(
+    query: torch.Tensor,
+    gallery: torch.Tensor,
+    k: int,
+    *,
+    exclude_idx: Optional[int] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Return ``(indices, scores)`` of the top-``k`` cosine neighbours of ``query``.
+
+    Args:
+        query: ``[d]`` or ``[1, d]`` dense query embedding.
+        gallery: ``[N, d]`` corpus of dense embeddings.
+        exclude_idx: optional gallery index to exclude (e.g. the seed track itself).
+    """
+    query = F.normalize(query.reshape(1, -1), dim=-1).squeeze(0)
+    gallery = F.normalize(gallery, dim=-1)
+    sims = gallery @ query
+    if exclude_idx is not None:
+        sims[exclude_idx] = -1e9
+    n_avail = gallery.shape[0] - (1 if exclude_idx is not None else 0)
+    k = min(int(k), int(n_avail))
+    vals, idx = torch.topk(sims, k=k, largest=True, sorted=True)
+    return idx.detach().cpu(), vals.detach().cpu()

steerable_retrieval/steer/slider.py

@@ -0,0 +1,245 @@
+"""The :class:`Slider` -- the public entry point for sparse steerable retrieval.
+
+A ``Slider`` binds a free-form text concept to a sparse edit direction in a trained
+SAE, obtained by sparse inversion (:mod:`steerable_retrieval.steer.inversion`). Once
+built it can steer any query embedding along the concept axis and retrieve over a
+corpus::
+
+    slider = Slider("distorted guitar", model=model)   # uses the packaged manifold prior
+    z_edited = slider.steer(z, alpha=1.0)               # alpha < 0 suppresses
+    idx, scores = slider.retrieve(z, corpus, alpha=1.0, k=10)
+
+The Mahalanobis **manifold prior** is a fixed distributional constant shipped with the
+library (fit on Music4All MuQ-MuLan embeddings) -- it is *independent of the retrieval
+corpus*. To adapt it to your own data, call :meth:`Slider.fit_prior` explicitly.
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+from steerable_retrieval.steer.inversion import (
+    InversionResult,
+    MahalanobisPrior,
+    fit_mahalanobis_prior,
+    invert_concept,
+    load_default_prior,
+)
+from steerable_retrieval.steer.retrieval import topk_cosine_neighbors
+from steerable_retrieval.steer.steering import build_edit_mask, decode_normalized, steer_sparse
+
+# Filled once the main MuQ/music4all SAE is trained and uploaded to HuggingFace.
+DEFAULT_MODEL_CLASS: Optional[str] = None
+DEFAULT_CHECKPOINT: Optional[str] = None
+# Name of the packaged manifold prior loaded when no prior is supplied.
+DEFAULT_PRIOR_NAME = "muq_mulan_music4all"
+
+ArrayLike = Union[np.ndarray, torch.Tensor]
+
+
+def _resolve_model(model, model_class, checkpoint):
+    if model is not None:
+        return model
+    if model_class is None and checkpoint is None and DEFAULT_CHECKPOINT is None:
+        raise ValueError(
+            "No model provided and no default checkpoint is available yet. "
+            "Pass an explicit `model=` (a trained SAE with .sae_encoder/.sae_decoder/"
+            ".text_encoder), or wait for the released pretrained checkpoint."
+        )
+    from steerable_retrieval.steer.loading import load_steerable_sae
+
+    return load_steerable_sae(
+        checkpoint or DEFAULT_CHECKPOINT, model_class=model_class or DEFAULT_MODEL_CLASS
+    )
+
+
+class Slider:
+    """A concept slider over a trained SAE's sparse feature space."""
+
+    def __init__(
+        self,
+        concept: str,
+        *,
+        model=None,
+        model_class: Optional[str] = DEFAULT_MODEL_CLASS,
+        checkpoint: Optional[str] = DEFAULT_CHECKPOINT,
+        prior: Optional[MahalanobisPrior] = None,
+        audio_embeddings: Optional[ArrayLike] = None,
+        method: str = "adam",
+        K: Optional[int] = None,
+        idf: Optional[torch.Tensor] = None,
+        norm_bundle: str = "norm",
+        device: Optional[Union[str, torch.device]] = None,
+        activation_threshold: float = 1e-6,
+        prior_name: str = DEFAULT_PRIOR_NAME,
+        **inversion_kwargs,
+    ):
+        """Build a slider for ``concept``.
+
+        Args:
+            concept: free-form text concept, e.g. ``"distorted guitar"``.
+            model: a trained SAE exposing ``sae_encoder``, ``sae_decoder``, ``text_encoder``.
+            model_class, checkpoint: alternative to ``model`` -- resolve a pretrained SAE.
+            prior: an explicit :class:`MahalanobisPrior`. If omitted, the packaged
+                default prior (``prior_name``) is used -- a fixed manifold constant,
+                *not* fit from any corpus. Use :meth:`fit_prior` to adapt it.
+            audio_embeddings: optional ``[N, d]`` audio embeddings used *only* to
+                initialize Adam inversion from the nearest audio neighbour. They do
+                **not** set the prior.
+            method: ``"adam"`` (default) or ``"fista"``.
+            K: cap the concept support to the top-``K`` features (``None`` = all active).
+            idf: optional IDF weights over dictionary features (down-weights hub neurons).
+            norm_bundle: edit-direction normalization (``"norm"`` = unit L2, recommended).
+            prior_name: which packaged prior to load when ``prior`` is omitted.
+        """
+        self.concept = concept
+        self.method = method
+        self.K = K
+        self.model = _resolve_model(model, model_class, checkpoint)
+        self.device = torch.device(device) if device is not None else self._infer_device()
+        self._idf = idf
+        self._norm_bundle = norm_bundle
+        self._activation_threshold = activation_threshold
+        self._inversion_kwargs = inversion_kwargs
+        self._audio_np = _to_numpy(audio_embeddings) if audio_embeddings is not None else None
+
+        # Prior precedence: explicit > packaged default constant. (audio_embeddings
+        # is an init aid only and never becomes the prior -- the prior is decoupled
+        # from any corpus. Call fit_prior() to adapt it on purpose.)
+        if prior is not None:
+            self.prior = prior.to(device=self.device)
+        else:
+            self.prior = load_default_prior(prior_name, device=self.device)
+        self._check_prior_dim()
+
+        self._text_embedding = self._embed_text(concept)
+        self._recompute()
+
+    # -- prior --------------------------------------------------------------- #
+    def _check_prior_dim(self) -> None:
+        d_model = int(getattr(self.model.sae_decoder, "act_size", self.model.sae_decoder.W_dec.shape[-1]))
+        d_prior = int(self.prior.mean.shape[-1])
+        if d_prior != d_model:
+            raise ValueError(
+                f"Manifold prior dim ({d_prior}) != model embedding dim ({d_model}). "
+                "Pass a matching `prior=`, fit one via fit_prior(embeddings), or set the "
+                "correct `prior_name`."
+            )
+
+    def fit_prior(self, embeddings: ArrayLike) -> "Slider":
+        """Refit the manifold prior to ``embeddings`` and recompute the slider.
+
+        Opt-in coupling: by default the slider uses the packaged prior, which is
+        independent of any corpus. Pass your retrieval corpus (or any representative
+        set) here to adapt the manifold prior to your own data. Returns ``self``.
+        """
+        self.prior = fit_mahalanobis_prior(_to_numpy(embeddings), device=self.device)
+        self._check_prior_dim()
+        self._recompute()
+        return self
+
+    def _recompute(self) -> None:
+        """Re-run inversion + rebuild the edit mask (after a prior/param change)."""
+        self.inversion: InversionResult = invert_concept(
+            self.model,
+            self._text_embedding,
+            self.prior,
+            method=self.method,
+            audio_embeddings=self._audio_np,
+            activation_threshold=self._activation_threshold,
+            **self._inversion_kwargs,
+        )
+        self.mask = build_edit_mask(
+            self.inversion.sparse_values,
+            dict_size=self.model.sae_encoder.dict_size,
+            K=self.K,
+            idf=self._idf,
+            norm_bundle=self._norm_bundle,
+            activation_threshold=self._activation_threshold,
+        )
+
+    # -- introspection ------------------------------------------------------- #
+    @property
+    def support(self) -> torch.Tensor:
+        """Indices of the active concept features (the slider's support)."""
+        return torch.nonzero(self.mask.abs() > 0, as_tuple=False).flatten()
+
+    def __len__(self) -> int:
+        return int(self.support.numel())
+
+    def __repr__(self) -> str:
+        return (
+            f"Slider(concept={self.concept!r}, method={self.method!r}, "
+            f"|support|={len(self)}, text_cos={self.inversion.final_text_cosine:.3f})"
+        )
+
+    # -- steering ------------------------------------------------------------ #
+    def steer(self, z: ArrayLike, alpha: float = 1.0) -> torch.Tensor:
+        """Steer dense query embedding(s) ``z`` along the concept axis.
+
+        ``alpha > 0`` amplifies the concept, ``alpha < 0`` suppresses it. Returns the
+        L2-normalized edited dense embedding(s) with the same leading shape as ``z``.
+        """
+        sparse, squeeze = self._encode_sparse(z)
+        mask = self.mask.to(device=sparse.device, dtype=sparse.dtype)
+        edited = steer_sparse(sparse, mask, alpha)
+        dense = decode_normalized(self.model, edited)
+        return dense.squeeze(0) if squeeze else dense
+
+    def amplify(self, z: ArrayLike, alpha: float = 1.0) -> torch.Tensor:
+        return self.steer(z, abs(alpha))
+
+    def suppress(self, z: ArrayLike, alpha: float = 1.0) -> torch.Tensor:
+        return self.steer(z, -abs(alpha))
+
+    # -- retrieval ----------------------------------------------------------- #
+    def retrieve(
+        self,
+        z: ArrayLike,
+        corpus: ArrayLike,
+        *,
+        alpha: float = 1.0,
+        k: int = 10,
+        exclude_idx: Optional[int] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Steer ``z`` by ``alpha`` then return top-``k`` cosine neighbours in ``corpus``.
+
+        ``corpus`` is an arbitrary set of dense embeddings to search over -- it is
+        independent of the manifold prior. Returns ``(indices, scores)`` into ``corpus``.
+        """
+        edited = self.steer(z, alpha)
+        gallery = _to_tensor(corpus, self.device)
+        return topk_cosine_neighbors(edited.to(self.device), gallery, k, exclude_idx=exclude_idx)
+
+    # -- internals ----------------------------------------------------------- #
+    def _infer_device(self) -> torch.device:
+        return self.model.sae_decoder.W_dec.device
+
+    def _embed_text(self, concept: str) -> torch.Tensor:
+        with torch.inference_mode():
+            emb = self.model.text_encoder([concept])
+        return emb.detach().clone().to(self.device)
+
+    def _encode_sparse(self, z: ArrayLike) -> Tuple[torch.Tensor, bool]:
+        t = _to_tensor(z, self.device)
+        squeeze = t.dim() == 1
+        if squeeze:
+            t = t.unsqueeze(0)
+        with torch.inference_mode():
+            _, sparse, _, _ = self.model.inference(t)
+        return sparse.clone(), squeeze
+
+
+def _to_numpy(x: ArrayLike) -> np.ndarray:
+    if isinstance(x, torch.Tensor):
+        return x.detach().cpu().numpy().astype(np.float32)
+    return np.asarray(x, dtype=np.float32)
+
+
+def _to_tensor(x: ArrayLike, device) -> torch.Tensor:
+    if isinstance(x, torch.Tensor):
+        return x.to(device=device, dtype=torch.float32)
+    return torch.from_numpy(np.asarray(x, dtype=np.float32)).to(device)

steerable_retrieval/steer/steering.py

@@ -0,0 +1,98 @@
+"""Steering operations: turn an inverted concept support into an edit direction.
+
+The concept support (from :mod:`steerable_retrieval.steer.inversion`) is converted to a masked,
+optionally IDF-reweighted, L2-normalized direction over the SAE dictionary. Steering
+adds (amplify) or subtracts (suppress) a scaled copy of that direction from a query's
+sparse code, then decodes back to the dense space for retrieval.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+
+
+def idf_weights(sparse_codes: torch.Tensor, *, eps: float = 1e-6, thresh: float = 1e-6) -> torch.Tensor:
+    """IDF weights over dictionary features from a corpus of sparse codes.
+
+    ``w_k = log(N / (df_k + eps) + 1)`` where ``df_k`` is the number of documents
+    (audio samples) in which feature ``k`` is active. Down-weights generic hub neurons.
+    """
+    num_docs = sparse_codes.shape[0]
+    df = (sparse_codes.abs() > thresh).float().sum(dim=0)
+    return torch.log(torch.tensor(float(num_docs)) / (df + eps) + 1.0)
+
+
+def apply_idf(values: torch.Tensor, idf: Optional[torch.Tensor]) -> torch.Tensor:
+    values = values.detach().cpu().float().view(-1)
+    if idf is not None and idf.shape[0] == values.shape[0]:
+        values = values * idf.detach().cpu().float().view(-1)
+    return values
+
+
+def normalize_bundle(weights: torch.Tensor, mode: str = "norm", eps: float = 1e-6) -> torch.Tensor:
+    """Normalize support weights so ``alpha`` is the sole edit-strength knob.
+
+    ``mode="norm"`` (default) makes the edit a unit-L2 direction, so support size
+    controls *which* coordinates move and ``alpha`` controls *how far*.
+    """
+    weights = weights.detach().cpu().float()
+    if not mode or weights.numel() == 0:
+        return weights
+    if mode == "prob":
+        total = float(weights.abs().sum().item())
+        return weights if total <= eps else weights / total
+    if mode == "normsum":
+        total = float(weights.abs().sum().item())
+        return weights / (total + eps) * float(weights.numel())
+    weight_norm = float(weights.norm().item())
+    return weights if weight_norm <= eps else weights / weight_norm
+
+
+def build_edit_mask(
+    sparse_values: torch.Tensor,
+    dict_size: int,
+    *,
+    K: Optional[int] = None,
+    idf: Optional[torch.Tensor] = None,
+    norm_bundle: str = "norm",
+    activation_threshold: float = 1e-6,
+) -> torch.Tensor:
+    """Build a dense ``[dict_size]`` edit direction from recovered sparse values.
+
+    Selects the top-``K`` active features (by |value|, after optional IDF reweighting),
+    normalizes them, and scatters them into a full-width mask.
+    """
+    values = apply_idf(sparse_values, idf)
+    active = torch.where(values.abs() > activation_threshold)[0]
+    if active.numel() == 0:
+        return torch.zeros(dict_size, dtype=torch.float32)
+    active_vals = values[active]
+    k = active.numel() if K is None else min(int(K), int(active.numel()))
+    order = torch.topk(active_vals.abs(), k=k).indices
+    idx = active[order]
+    vals = normalize_bundle(active_vals[order], mode=norm_bundle)
+    mask = torch.zeros(dict_size, dtype=torch.float32)
+    mask[idx.long()] = vals.float()
+    return mask
+
+
+def steer_sparse(sparse: torch.Tensor, mask: torch.Tensor, alpha: float, *, clamp_min_zero: bool = True) -> torch.Tensor:
+    """Steer a sparse code along the edit direction.
+
+    ``alpha > 0`` amplifies the concept, ``alpha < 0`` suppresses it. Suppression
+    clamps at zero to respect ReLU-style non-negative SAE codes.
+    """
+    mask = mask.to(device=sparse.device, dtype=sparse.dtype)
+    new_sparse = sparse + alpha * mask
+    if clamp_min_zero and alpha < 0:
+        new_sparse = new_sparse.clamp_min(0.0)
+    return new_sparse
+
+
+def decode_normalized(model, sparse: torch.Tensor) -> torch.Tensor:
+    """Decode a sparse code to the dense space and L2-normalize (retrieval geometry)."""
+    dense = model.sae_decoder(sparse)
+    return F.normalize(dense, dim=-1)

steerable_retrieval/train.py

@@ -0,0 +1,243 @@
+import os
+from typing import Any, Dict, List, Optional, Tuple
+
+from dora import get_xp, hydra_main
+import hydra
+
+import lightning as L
+import rootutils
+import torch
+from lightning import Callback, LightningDataModule, LightningModule, Trainer
+from lightning.pytorch.loggers import Logger
+from omegaconf import DictConfig
+import logging
+from pathlib import Path
+
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+# ------------------------------------------------------------------------------------ #
+# the setup_root above is equivalent to:
+# - adding project root dir to PYTHONPATH
+#       (so you don't need to force user to install project as a package)
+#       (necessary before importing any local modules e.g. `from gdr import utils`)
+# - setting up PROJECT_ROOT environment variable
+#       (which is used as a base for paths in "configs/paths/default.yaml")
+#       (this way all filepaths are the same no matter where you run the code)
+# - loading environment variables from ".env" in root dir
+#
+# you can remove it if you:
+# 1. either install project as a package or move entry files to project root dir
+# 2. set `root_dir` to "." in "configs/paths/default.yaml"
+#
+# more info: https://github.com/ashleve/rootutils
+# ------------------------------------------------------------------------------------ #
+
+from steerable_retrieval.utils import (
+    RankedLogger,
+    extras,
+    get_metric_value,
+    instantiate_callbacks,
+    instantiate_loggers,
+    log_hyperparameters,
+    register_resolvers,
+    task_wrapper,
+)
+
+log = RankedLogger(__name__, rank_zero_only=True)
+register_resolvers()
+
+
+@task_wrapper
+def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+    """Trains the model. Can additionally evaluate on a testset, using best weights obtained during
+    training.
+
+    This method is wrapped in optional @task_wrapper decorator, that controls the behavior during
+    failure. Useful for multiruns, saving info about the crash, etc.
+
+    :param cfg: A DictConfig configuration composed by Hydra.
+    :return: A tuple with metrics and dict with all instantiated objects.
+    """
+    # set seed for random number generators in pytorch, numpy and python.random
+    if cfg.get("seed"):
+        L.seed_everything(cfg.seed, workers=True)
+
+    log.info(f"Instantiating datamodule <{cfg.data._target_}>")
+    datamodule = hydra.utils.instantiate(cfg.data)
+
+    log.info(f"Instantiating model <{cfg.model._target_}>")
+    model: LightningModule = hydra.utils.instantiate(cfg.model)
+    # model.xp = get_xp()
+
+    log.info("Instantiating callbacks...")
+    callbacks: List[Callback] = instantiate_callbacks(cfg.get("callbacks"))
+    log.info(f"Callbacks: {callbacks}")
+
+    log.info("Instantiating loggers...")
+    logger: List[Logger] = instantiate_loggers(cfg.get("logger"))
+
+    log.info(f"Instantiating trainer <{cfg.trainer._target_}>")
+    trainer: Trainer = hydra.utils.instantiate(cfg.trainer, logger=logger, callbacks=callbacks)
+
+    object_dict = {
+        "cfg": cfg,
+        "datamodule": datamodule,
+        "model": model,
+        "callbacks": callbacks,
+        "logger": logger,
+        "trainer": trainer,
+    }
+
+    if logger:
+        log.info("Logging hyperparameters!")
+        log_hyperparameters(object_dict)
+
+    # automatically resume from latest checkpoint if exists and ckpt_path not manually specified
+    # TODO: discuss cfg.resume, this is anti-dora but maybe it's useful
+    ckpt_path = cfg.get("ckpt_path")
+    cfg.resume = cfg.resume or os.environ.get("USE_MPI")
+
+    if '/opt/ml/' in cfg.paths.ckpt_dir:
+        was_s3 = True
+    else:
+        was_s3 = False
+
+
+    logging.info("="*100)
+    # logging.info(os.listdir('/opt/ml/input/data')) if os.path.exists('/opt/ml/input/data') else logging.info("No data found in /opt/ml/input/data")
+    # log tree of /opt/ml/input/data
+
+    def tree_str(
+        path=".",
+        max_depth=None,
+        max_files=2,
+        ignore={".git", "__pycache__"}
+    ):
+        lines = []
+        path = Path(path)
+
+        def _walk(p, prefix="", level=0):
+            if max_depth is not None and level > max_depth:
+                return
+
+            entries = [e for e in p.iterdir() if e.name not in ignore]
+
+            dirs = sorted((e for e in entries if e.is_dir()), key=lambda x: x.name.lower())
+            files = sorted((e for e in entries if e.is_file()), key=lambda x: x.name.lower())
+
+            shown_files = files[:max_files]
+            omitted_files = len(files) - len(shown_files)
+
+            combined = dirs + shown_files
+
+            for i, entry in enumerate(combined):
+                is_last = i == len(combined) - 1
+                connector = "└── " if is_last else "├── "
+                lines.append(prefix + connector + entry.name)
+
+                if entry.is_dir():
+                    extension = "    " if is_last else "│   "
+                    _walk(entry, prefix + extension, level + 1)
+
+            if omitted_files > 0:
+                lines.append(prefix + f"└── … ({omitted_files} more files)")
+
+        _walk(path)
+        return "\n".join(lines)
+
+    logging.info("="*100)
+    data_root = cfg.paths.get('data_dir')  # config-provided data path (was hardcoded /opt/ml on SageMaker)
+    if data_root and os.path.exists(data_root):
+        logging.info(tree_str(data_root))
+    logging.info("="*100)
+
+    if os.path.exists(cfg.paths.ckpt_dir) and cfg.resume:
+        candidates = [os.path.join(cfg.paths.ckpt_dir, ckpt_file) for ckpt_file in os.listdir(cfg.paths.ckpt_dir) if ckpt_file.endswith(".ckpt")]
+        if candidates:
+            # get the last modified ckpt else get last.ckpt, reason is that s3 downloads are not in order of creation
+            # ckpt_path = max(candidates, key=os.path.getmtime) if not was_s3 else 
+
+            if was_s3:
+                ckpt_path = os.path.join(cfg.paths.ckpt_dir, "last.ckpt")
+                if "last.ckpt" not in os.listdir(cfg.paths.ckpt_dir):
+                    log.warning("last.ckpt not found in s3 ckpt_dir. Training from scratch!")
+                    ckpt_path = None
+            else:
+                ckpt_path = max(candidates, key=os.path.getmtime)
+                log.info(f"Resuming from checkpoint {ckpt_path}...")
+
+            # ckpt_path = os.path.join(cfg.paths.ckpt_dir, "last.ckpt") if "last.ckpt" in os.listdir(cfg.paths.ckpt_dir) else None
+            log.info(f"Resuming from checkpoint {ckpt_path}...")
+        else:
+            log.info(ckpt_path, "is empty. Training from scratch!")
+
+    
+    trainer.true_accumulate_grad_batches, trainer.accumulate_grad_batches = trainer.accumulate_grad_batches, 1
+    model.gradient_clip_val, trainer.gradient_clip_val = trainer.gradient_clip_val, None
+
+    if cfg.get("train"):
+        log.info("Starting training!")
+        with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=True):
+            trainer.fit(model=model, datamodule=datamodule, ckpt_path=ckpt_path)
+
+    train_metrics = trainer.callback_metrics
+
+    if cfg.get("test"):
+        log.info("Starting testing!")
+        # Get best checkpoint path if checkpoint callback exists
+        if hasattr(trainer, 'checkpoint_callback') and trainer.checkpoint_callback is not None:
+            ckpt_path = trainer.checkpoint_callback.best_model_path
+            if ckpt_path == "":
+                log.warning("Best ckpt not found! Using current weights for testing...")
+                ckpt_path = None
+        else:
+            log.warning("No checkpoint callback found! Using current weights for testing...")
+            ckpt_path = None
+        trainer.test(model=model, datamodule=datamodule, ckpt_path=ckpt_path)
+        log.info(f"Best ckpt path: {ckpt_path}")
+
+    test_metrics = trainer.callback_metrics
+
+    # merge train and test metrics
+    metric_dict = {**train_metrics, **test_metrics}
+
+    return metric_dict, object_dict
+
+    return {}, object_dict
+
+
+@hydra_main(version_base="1.3", config_path="../configs", config_name="train.yaml")
+def main(cfg: DictConfig) -> Optional[float]:
+    """Main entry point for training.
+
+    :param cfg: DictConfig configuration composed by Hydra.
+    :return: Optional[float] with optimized metric value.
+    """
+    # handle A100 GPUs
+    if torch.cuda.is_available() and ("A100" in torch.cuda.get_device_name() or "A5000" in torch.cuda.get_device_name()):
+        torch.set_float32_matmul_precision("high")
+
+    # avoid annoying multiprocessing errors
+    torch.multiprocessing.set_sharing_strategy('file_system')
+
+    # prevent annoying warning
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+    # apply extra utilities
+    # (e.g. ask for tags if none are provided in cfg, print cfg tree, etc.)
+    extras(cfg)
+
+    # train the model
+    metric_dict, _ = train(cfg)
+
+    # safely retrieve metric value for hydra-based hyperparameter optimization
+    metric_value = get_metric_value(
+        metric_dict=metric_dict, metric_name=cfg.get("optimized_metric")
+    )
+
+    # return optimized metric
+    return metric_value
+
+
+if __name__ == "__main__":
+    main()

steerable_retrieval/utils/__init__.py

@@ -0,0 +1,6 @@
+from steerable_retrieval.utils.instantiators import instantiate_callbacks, instantiate_loggers
+from steerable_retrieval.utils.logging_utils import log_hyperparameters
+from steerable_retrieval.utils.pylogger import RankedLogger
+from steerable_retrieval.utils.resolvers import register_resolvers
+from steerable_retrieval.utils.rich_utils import enforce_tags, print_config_tree
+from steerable_retrieval.utils.utils import extras, get_metric_value, task_wrapper

steerable_retrieval/utils/copy.py

@@ -0,0 +1,67 @@
+r"""Utils for copying data to the temporary directory of the compute node
+"""
+
+import logging
+import os
+from pathlib import Path
+
+from filelock import FileLock
+
+
+log = logging.getLogger(__name__)
+
+
+def copy_data(origin: Path, destination: Path):
+    cmd = "cp -r {origin} {destination}"
+    os.system(cmd.format(origin=origin, destination=destination))
+
+def copy_to_compute_node(data_path: str | Path,
+                         local_dir: str | Path = "/local/job"):
+    if not isinstance(data_path, Path):
+        data_path = Path(data_path)
+    
+    if not isinstance(local_dir, Path):
+        local_dir = Path(local_dir)
+    
+    # If we are not in a SLURM job, this function is a no-op
+    job_id = os.environ.get("SLURM_JOB_ID")
+    if job_id is None:
+        return data_path
+    
+    # First, we check if the data we're looking for is already in the compute node
+    origin_folder = None
+
+    for folder in local_dir.iterdir():
+        if not os.access(folder, os.R_OK):
+            continue
+
+        subfolder = folder / data_path.name
+        if subfolder.exists():
+            origin_folder = subfolder
+            break
+
+    # Set the destination folder
+    dest_folder = local_dir / job_id / data_path.name
+    if dest_folder.exists():
+        log.info(f"Data found in {dest_folder}.")
+        return dest_folder
+    
+    dest_folder.mkdir()
+
+    if origin_folder is None:
+        # In that case, we will copy from the login node directly (slow...)
+        origin_folder = data_path
+    
+    else:
+        # Wait for the origin folder to have been fully filled to copy from the compute node (much faster!)
+        log.info(f"Data found in {origin_folder}. Waiting for the lock to be released...")
+        with FileLock(origin_folder / "lock"):
+            pass
+    
+    # Lock the destination folder and copy the data inside it
+    log.info(f"Copying data from {origin_folder} to {dest_folder}...")
+    with FileLock(dest_folder / "lock"):
+        copy_data(origin_folder, dest_folder)
+    
+    log.info("Done.")
+    return dest_folder