""" 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()