sae_backend.py

"""
SAE-Based Steering Backend for Gradio App.

Provides functions to load a pretrained Sparse Autoencoder (MSAE),
encode/decode CLIP embeddings through the SAE latent space, and
perform Pseudo-Relevance Feedback (PRF) steering.

Adapted from sae_experiments/sae_utils.py for self-contained deployment.
"""

import numpy as np
import torch
from pathlib import Path

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

SAE_HF_REPO = "WolodjaZ/MSAE"
SAE_WEIGHTS_PATH = (
    "ViT-B_16/centered/4096_512_TopKReLU_64_RW_False_False_0.0"
    "_cc3m_ViT-B~16_train_image_2905936_512.pth"
)
CONCEPT_INTERPRETER_PATH = (
    "ViT-B_16/centered/Concept_Interpreter_4096_512_TopKReLU_64_RW"
    "_False_False_0.0_cc3m_ViT-B~16_train_image_2905936_512"
    "_disect_ViT-B~16_-1_text_20000_512.npy"
)

VOCAB_FILE = Path(__file__).parent / "vocab" / "clip_disect_20k.txt"
MEAN_ACTIVATIONS_FILE = Path(__file__).parent / "mean_activations.pt"

# ---------------------------------------------------------------------------
# Global cache
# ---------------------------------------------------------------------------

_sae_model = None
_mean_activations = None

# ---------------------------------------------------------------------------
# SAE Loading
# ---------------------------------------------------------------------------


def _download_sae_weights():
    """Download pretrained SAE weights from HuggingFace Hub.

    Returns:
        (local_weights_path, local_concept_path)
    """
    from huggingface_hub import hf_hub_download

    weights_path = hf_hub_download(repo_id=SAE_HF_REPO, filename=SAE_WEIGHTS_PATH)
    concept_path = hf_hub_download(repo_id=SAE_HF_REPO, filename=CONCEPT_INTERPRETER_PATH)
    return weights_path, concept_path


def load_sae_model():
    """Load the pretrained SAE (cached after first call).

    Returns:
        SAE model on CPU (or CUDA if available).
    """
    global _sae_model
    if _sae_model is not None:
        return _sae_model

    from msae.sae import SAE

    print("🔄 Downloading SAE weights from HuggingFace Hub…")
    weights_path, _ = _download_sae_weights()

    # MSAE load_model parses config from filename using path.split("/"),
    # which breaks on Windows backslash paths.  Normalise to forward slashes.
    weights_path = str(weights_path).replace("\\", "/")

    # PyTorch 2.6+ defaults weights_only=True which rejects numpy scalars
    # in older checkpoints.  Allow them before loading.
    torch.serialization.add_safe_globals(
        [np.core.multiarray.scalar, np.dtype, np.dtypes.Float64DType]
    )

    sae_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    _sae_model = SAE(weights_path)
    _sae_model = _sae_model.to(sae_device)
    _sae_model.eval()
    print(f"✅ SAE loaded ({_sae_model.model.n_latents} latents) on {sae_device}")
    return _sae_model


def load_mean_activations():
    """Load dataset-level mean SAE activations (cached).

    Returns:
        torch.Tensor of shape (n_latents,)
    """
    global _mean_activations
    if _mean_activations is not None:
        return _mean_activations

    if MEAN_ACTIVATIONS_FILE.exists():
        _mean_activations = torch.load(
            str(MEAN_ACTIVATIONS_FILE), map_location="cpu", weights_only=True
        )
        print("✅ Loaded mean_activations.pt")
    else:
        print("⚠️ mean_activations.pt not found — PRF steering will work without centering")
        _mean_activations = None
    return _mean_activations


# ---------------------------------------------------------------------------
# SAE Encode / Decode
# ---------------------------------------------------------------------------


def sae_encode(sae_model, embeddings):
    """Encode CLIP embeddings through the SAE.

    Args:
        sae_model: loaded SAE model
        embeddings: np.ndarray (N, 512) or torch.Tensor

    Returns:
        latents: torch.Tensor (N, 4096) on CPU
    """
    if isinstance(embeddings, np.ndarray):
        embeddings = torch.from_numpy(embeddings).float()
    sae_device = sae_model.mean.device
    embeddings = embeddings.to(sae_device)
    with torch.no_grad():
        _, full_latents = sae_model.encode(embeddings)
    return full_latents.detach().cpu()


def sae_decode(sae_model, latents):
    """Decode SAE latents back to CLIP embedding space.

    Args:
        sae_model: loaded SAE model
        latents: torch.Tensor (N, n_latents)

    Returns:
        reconstructed: np.ndarray (N, 512)
    """
    sae_device = sae_model.mean.device
    if isinstance(latents, np.ndarray):
        latents = torch.from_numpy(latents).float()
    latents = latents.to(sae_device)
    with torch.no_grad():
        reconstructed = sae_model.model.decode(latents)
        reconstructed = sae_model.postprocess(reconstructed)
    return reconstructed.cpu().numpy().astype(np.float32)


# ---------------------------------------------------------------------------
# PRF Steering
# ---------------------------------------------------------------------------


def sae_prf_steering(
    query_emb,
    image_embeddings,
    baseline_indices,
    *,
    prf_k=10,
    top_feats=32,
    scale=1.0,
):
    """Pseudo-Relevance Feedback (PRF) steering in SAE latent space.

    Uses the baseline CLIP retrieval top-K images as pseudo-positives,
    extracts their dominant SAE feature residuals, then nudges the query's
    SAE latents along those residual directions before decoding back to
    CLIP embedding space.

    Args:
        query_emb: np.ndarray (512,) or (1, 512)  — query CLIP embedding
        image_embeddings: np.ndarray (N, 512)      — dataset image embeddings
        baseline_indices: np.ndarray/list           — baseline top-K image indices
        prf_k: number of baseline images to use as pseudo-positives
        top_feats: number of SAE features to modify (by |mean residual|)
        scale: step size applied to mean residual on selected features

    Returns:
        steered_emb: np.ndarray (512,)  — L2-normalised steered embedding
    """
    sae_model = load_sae_model()
    mean_acts = load_mean_activations()

    if isinstance(baseline_indices, np.ndarray):
        idx = baseline_indices
    else:
        idx = np.array(list(baseline_indices), dtype=np.int64)

    prf_k = int(min(prf_k, len(idx)))
    q = query_emb.reshape(1, -1)

    if prf_k <= 0:
        # No pseudo-positives available — return normalised query
        q_norm = q / (np.linalg.norm(q, axis=-1, keepdims=True) + 1e-8)
        return q_norm.flatten()

    # Encode query and top-K images through SAE
    q_lat = sae_encode(sae_model, q)          # (1, D)
    top_embs = image_embeddings[idx[:prf_k]]
    top_lat = sae_encode(sae_model, top_embs)  # (K, D)

    # Centre w.r.t. mean activations if available
    if mean_acts is not None:
        mean = mean_acts.detach().cpu().unsqueeze(0)  # (1, D)
        q_resid = q_lat - mean
        top_resid = top_lat - mean
    else:
        mean = None
        q_resid = q_lat
        top_resid = top_lat

    # Average residual across pseudo-positives
    mean_resid = top_resid.mean(dim=0)  # (D,)
    d = int(mean_resid.numel())
    k = int(min(top_feats, d))

    if k <= 0:
        q_norm = q / (np.linalg.norm(q, axis=-1, keepdims=True) + 1e-8)
        return q_norm.flatten()

    # Select top features by |mean residual|
    selected = torch.topk(mean_resid.abs(), k=k).indices

    # Nudge query residuals along those features
    q_resid = q_resid.clone()
    q_resid[0, selected] = q_resid[0, selected] + float(scale) * mean_resid[selected]

    # Reconstruct full latent and decode
    if mean is not None:
        steered_lat = q_resid + mean
    else:
        steered_lat = q_resid

    steered = sae_decode(sae_model, steered_lat)  # (1, 512)

    # L2-normalise for cosine retrieval
    steered = steered / (np.linalg.norm(steered, axis=-1, keepdims=True) + 1e-8)
    return steered.flatten()