sae/SAE_Tools.py

import torch as t
from torch import Tensor
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import einops
from functools import partial
from typing import List, Tuple, Dict, Any, Union, Callable, cast, Literal
from tqdm import tqdm
from collections import OrderedDict
from transformer_lens import (
    HookedTransformer,
    ActivationCache,
)
import math
from .autoencoder import TopKSAE, SAEOut, SAE, SAE_Template, TopKTranscoder, TopKCrosscoder, BatchTopKSAE
from transformer_lens.hook_points import HookedRootModule, HookPoint
from sae_lens import SAE as SAE_LENS
from model.blip.configuration_blip import BlipConfig
from model.llava.configuration_llava import LlavaConfig


from torch.utils.data import DataLoader, TensorDataset
from extra_materials.graph.Graph_Template import Node, Index

import plotly.express as px
from jaxtyping import Float, Int
from rich import print as rprint
from rich.table import Table
from tabulate import tabulate

def load_sae_model(file_path: str, model_type: Literal["llava", "blip"], hook_type: Literal["vision", "text"] = "text", device: str | t.device = "cpu", **kwargs) -> SAE_Template:
    '''
    Load TopKSAE from path
    '''
    tmp = t.load(file_path, map_location=device)
    
    sae_state_dict = {k: v for k, v in tmp["state_dict"].items() if "autoencoder" in k}
    sae_state_dict = {k.replace("autoencoder.", ""): v for k, v in sae_state_dict.items()}
      
    hyperparameters = tmp["hyper_parameters"]
    
    if model_type == "llava":
        model_cfg = LlavaConfig()
        if hook_type == "vision":
            hidden_size = model_cfg.vision_config.hidden_size
        elif hook_type == "text":
            hidden_size = model_cfg.text_config.hidden_size
        else:
            raise ValueError(f"Unknown hook type: {hook_type}")
    elif model_type == "blip":
        model_cfg = BlipConfig()
        if hook_type == "vision":
            hidden_size = model_cfg.vision_config.hidden_size
        elif hook_type == "text":
            hidden_size = model_cfg.text_config.hidden_size
        else:
            raise ValueError(f"Unknown hook type: {hook_type}")
    else: 
        raise ValueError(f"Unknown model type: {model_type}")
    
    
    dead_steps_threshold = hyperparameters["dead_steps_threshold"] or hyperparameters["dead_tokens_threshold"] // (
        hyperparameters["batch_size"] * hyperparameters["max_length"] * hyperparameters["accumulate_grad_batches"]
    )
    
    sae_name = hyperparameters.get("sae_name", "topk")
    
    if sae_name == "sae":
        class_sae = SAE
    elif sae_name == "topk":
        class_sae = TopKSAE
    elif sae_name == "batchtopk":
        class_sae = TopKSAE
    elif sae_name == "topk_transcoder":
        class_sae = TopKTranscoder
    else:
        raise ValueError(f"Unknown SAE name: {hyperparameters['sae_name']}")
        
    sae = class_sae(
        d_in=hidden_size,
        d_sae=int(hyperparameters["expansion_factor"] * hidden_size),
        hook_names=hyperparameters["hook_names"],
        k=hyperparameters["k"],
        dead_steps_threshold=dead_steps_threshold,
        dead_threshold=hyperparameters["dead_threshold"],
        auxk=hyperparameters["auxk"],
        standardize=hyperparameters["standardize"],
        **kwargs,
    ).to(device)
    sae.load_state_dict(sae_state_dict, strict=False)
    return sae

def load_crosscoder_model(file_path: str, model: HookedTransformer, device: str | t.device = "cpu", **kwargs) -> List[SAE_Template]:
    '''
    Load TopKSAE from path
    '''
    tmp = t.load(file_path, map_location=device)
      
    hyperparameters = tmp["hyper_parameters"]
    model_cfg = model.cfg
    dead_steps_threshold = hyperparameters["dead_steps_threshold"] or hyperparameters["dead_tokens_threshold"] // (
        hyperparameters["batch_size"] * hyperparameters["max_length"] * hyperparameters["accumulate_grad_batches"]
    )
    sae_name = hyperparameters.get("sae_name", "topk_crosscoder")
    
    if sae_name == "topk_crosscoder":
        class_sae = TopKCrosscoder
    else:
        raise ValueError(f"Unknown SAE name: {hyperparameters['sae_name']}")
    
    saes = []
    for layer in range(model.cfg.n_layers):
        sae_state_dict = {k: v for k, v in tmp["state_dict"].items() if f"crosscoders.{layer}." in k}
        sae_state_dict = {k.replace(f"crosscoders.{layer}.", ""): v for k, v in sae_state_dict.items()}
        sae = class_sae(
            d_in=model_cfg.d_model,
            d_sae=int(hyperparameters["expansion_factor"] * model_cfg.d_model),
            input_hook=hyperparameters["input_hooks"][layer],
            output_hooks=hyperparameters["output_hooks"][layer:],
            k=hyperparameters["k"],
            dead_steps_threshold=dead_steps_threshold,
            dead_threshold=hyperparameters["dead_threshold"],
            auxk=hyperparameters["auxk"],
            standardize=hyperparameters["standardize"],
            **kwargs,
        ).to(device)
        sae.load_state_dict(sae_state_dict)
        saes.append(sae)
        
    return saes

def normalize_activations(
    activations: Tensor,
    eps: float = 1e-6,
) ->  Tensor:
    '''
    Normalize activations
    '''
    exclude_first_dim = [i for i in range(1, activations.ndim)]
    mean = activations.mean(dim=exclude_first_dim, keepdim=True)
    std = activations.std(dim=exclude_first_dim, keepdim=True)
    activations = (activations - mean) / (std + eps)
    return activations

def get_sae_activations(
    sae: HookedRootModule,
    dataloader: DataLoader,
    device: str | t.device,
    normalize: bool = False,
    no_tqdm: bool = False,
) -> Tuple[Tensor, Tensor]:
    '''
    Return the reconstruction, indices and values of HookedRootModule given the cached activations dataset
    '''
    with t.no_grad():
        sae.eval()
        sae.to(device)
        
        if isinstance(sae, TopKSAE) or isinstance(sae, BatchTopKSAE):
            indices = []
            values = []
            for act in tqdm(dataloader, disable=no_tqdm):
                act = act[0]
                act = normalize_activations(act.to(device)) if normalize else act.to(device)
                out: SAEOut = sae.forward_training(act)
                indices.append(out.topk.indices.to("cpu"))
                values.append(out.topk.values.to("cpu"))
            indices = t.cat(indices)
            values = t.cat(values)
        else:
            latents_list: List[Tensor] = []
            max_k = [0] # to change outside the function
            def hook_fn(inputs: Tensor, hook: HookPoint, max_k = max_k) -> None:
                if "post" in hook.name:  # type: ignore
                    # (b, s, d)
                    inputs = inputs
                    mask = inputs > 0
                    max_k[0] = max(max_k[0], mask.sum(-1).max().item()) # update the max_k
                    
                    # convert to sparse tensor to save memory
                    indices = mask.nonzero(as_tuple=False).t() 
                    values = inputs[mask] 
                    sparse_tensor = t.sparse_coo_tensor(indices, values, inputs.shape).coalesce().detach().cpu()
                    latents_list.append(sparse_tensor)

            for act in tqdm(dataloader, disable=no_tqdm):
                act = act[0]
                act = normalize_activations(act.to(device)) if normalize else act.to(device)
                sae.run_with_hooks(
                    act, 
                    fwd_hooks=[(lambda name: "acts_post" in name or "recons" in name, hook_fn)],
                )

            indices_list = []
            values_list = []
            # transform into values and indices
            for sparse_latent in latents_list:
                latents = sparse_latent.to_dense() # (b, s, d)
                values, indices = t.topk(
                    latents, k=max_k[0], sorted=False,
                )
                indices_list.append(indices)
                values_list.append(values)
                
            indices = t.cat(indices_list)
            values = t.cat(values_list)
    
    return indices, values


def show_activation_histogram(
    sae: SAE_Template | SAE_LENS,
    values: Tensor,
    indices: Tensor,
    threshold: float = 1e-3,
) -> None:
    '''
    Show activation histogram given the values and indices of SAE_Template or SAE from sae_lens
    '''
    batch_size, pos, k = values.shape
    
    activation_counts = t.zeros(sae.cfg.d_sae, dtype=t.float32)
    activated = values > threshold
    activated_indices = indices[activated]

    # accumulate activation counts
    activation_counts.scatter_add_(0, activated_indices.flatten(), t.ones_like(activated_indices, dtype=t.float32).flatten())

    # Compute activation fraction
    activation_fraction = activation_counts / (batch_size * pos)
    

    px.histogram(
        activation_fraction.cpu().numpy(),
        nbins=50,
        title=f"ACTIVATIONS DENSITY",
        labels={"value": "Activation"},
        width=800,
        template="ggplot2",
        color_discrete_sequence=["darkorange"],
    ).update_layout(bargap=0.02, showlegend=False).show()


def show_feature_act_histogram(
    sae: SAE_Template | SAE_LENS,
    feat_idx: int,
    values: Tensor, 
    indices: Tensor, 
    threshold=1e-3,
) -> None:
    '''
    Show feature activation histogram given the values and indices of SAE_Template or SAE in sae_lens
    '''
    batch_size, pos, k = values.shape
    
    activation_counts = t.zeros(sae.cfg.d_sae, dtype=t.float32)
    activated = values > threshold
    # filter out non-activated indices
    activated_indices = t.where(activated, indices, t.zeros_like(indices))

    # filter out feature activations
    activation_feat = values[activated_indices.long() == feat_idx] 
    
    px.histogram(
        activation_feat.cpu().numpy(),
        nbins=50,
        title=f"ACTIVATIONS DENSITY",
        labels={"value": "Activation"},
        width=800,
        template="ggplot2",
        color_discrete_sequence=["darkorange"],
    ).update_layout(bargap=0.02, showlegend=False).show()
    
def cache_activation(
    transformer: HookedTransformer,
    input_toks: Tensor,
    hook_names: List[str] | None = None,
    batch_size: int = 256,
    important_components: List[Tuple[Node, Index]] | None = None,
    return_toks: bool = False,
    save_path: str | None = None,
    no_verbose: bool = False,
) -> Dict[str, Tensor] | Tuple[Dict[str, Tensor], Dict[str, Tensor]]:
    '''
    Caches the activations of the transformer for the given input tokens.
    '''
    if save_path is not None:
        try:
            all_cache = t.load(save_path)
            if not no_verbose:
                print("Activations loaded from cache.")
            return all_cache
        except:
            if not no_verbose:
                print("No cache found.")
            pass
    
    toks_loader = get_loader(input_toks, batch_size)
    
    if not no_verbose:
        print("Caching activations...")
    if important_components is not None:
        all_cache = OrderedDict({comp.name + repr(idx): [] for comp, idx in important_components})
        all_toks = OrderedDict({comp.name + repr(idx): [] for comp, idx in important_components})
        names_filter = lambda x: x in [comp.name for comp, _ in important_components]
    else:
        assert hook_names is not None, "Either important_components or hook_names be not None."
        all_cache = OrderedDict({name: [] for name in hook_names})
        all_toks = OrderedDict({name: [] for name in hook_names})
        names_filter = lambda x: x in hook_names
    
    for batch in tqdm(toks_loader, disable=no_verbose):
        batch_toks = batch[0]
        _, tmp_cache = transformer.run_with_cache(batch_toks, names_filter=names_filter)
        
        for name, cache in tmp_cache.items():
            if important_components is not None:
                for comp, idx in important_components:
                    if name == comp.name:
                        all_cache[name + repr(idx)].append(cache[idx.as_index].clone().cpu())
                        if return_toks:
                            all_toks[name + repr(idx)].append(batch_toks.clone().cpu())
            else:
                if name.endswith("result") and important_components is None:
                    cache = einops.rearrange(cache, "b s h d -> (b h) s d")
                all_cache[name].append(cache.cpu())
                if return_toks:
                    all_toks[name].append(batch_toks.clone().cpu())
    
    for tens in all_cache.values():
        assert len(tens) > 0, "No activations cached."
    
    all_cache = {name: t.cat(cache) for name, cache in all_cache.items()}
    if return_toks:
        all_toks = {name: t.cat(toks) for name, toks in all_toks.items()}
        return all_cache, all_toks
    if save_path is not None:
        t.save(all_cache, save_path)
    return all_cache    

def get_cache_with_toks_dataset(
    model: HookedTransformer,
    toks: Tensor, 
    hook_names: List[str] | None,
    batch_size: int,
    list_components: List[Tuple[Node, Index]] | None = None,
    no_verbose: bool = False
) -> Tuple[Dict[str, Tensor], Dict[str, Tensor]]:
    '''
    Cache the model activations given the toks and the list of component to cache.
    The return is a dictionary of activations and toks
    '''
    
    if toks.ndim == 1:
        # single input string
        toks = toks.unsqueeze(0)
    # shape must be [batch, ...]
    activation_cache, toks_cache = cache_activation(
        model,
        toks,
        hook_names,
        batch_size,
        list_components,
        return_toks=True,
        no_verbose=no_verbose,
    )
    
    return activation_cache, toks_cache # type: ignore

def get_loader_from_dict(list_dict: List[Dict[str, Tensor]], batch_size: int) -> DataLoader:
    '''
    Get loader from a dictionary of tensors
    '''
    list_tensor = []
    for d in list_dict:
        list_tensor.append(t.cat(list(d.values()), dim=0))
    
    dataset = TensorDataset(*list_tensor)
    return DataLoader(dataset, batch_size=batch_size, shuffle=False) # must be false, don't shuffle

def get_loader(toks: Tensor, batch_size: int) -> DataLoader:
    '''
    Get loader from toks
    '''
    if toks.ndim == 1:
        toks = toks.unsqueeze(0)
    dataset = TensorDataset(toks)
    return DataLoader(dataset, batch_size=batch_size, shuffle=False) # must be false, don't shuffle

def get_cache_with_toks_dataset_on_target_components(
    sae: SAE_Template | SAE_LENS,
    model: HookedTransformer,
    toks: Tensor,
    list_components: List[Tuple[Node, Index]],
    batch_size: int,
    device: str | t.device = "cpu",
    no_verbose: bool = False,
) -> Tuple[Tensor, Tensor, Dict[str, Tensor], Dict[str, Tensor]]:
    '''
    Get the activation cache and toks cache for the list components
    '''
    activation_cache, toks_cache = get_cache_with_toks_dataset(
        model, 
        toks, 
        [],
        batch_size,
        list_components, 
        no_verbose,
    )

    indices, values = get_sae_activations(
        sae, 
        get_loader_from_dict([activation_cache], batch_size), 
        device, 
        False, 
        no_verbose,
    )
    return indices, values, toks_cache, activation_cache

def get_cache_with_toks_dataset_on_target_hooks(
    sae: SAE_Template | SAE_LENS,
    model: HookedTransformer,
    toks: Tensor,
    hook_names: List[str],
    batch_size: int,
    device: str | t.device = "cpu",
    no_verbose: bool = False,
) -> Tuple[Tensor, Tensor, Dict[str, Tensor], Dict[str, Tensor]]:
    '''
    Get the activation cache and toks cache for the list hook_names
    '''
    activation_cache, toks_cache = get_cache_with_toks_dataset(
        model, 
        toks, 
        hook_names,
        batch_size,
        None, 
        no_verbose,
    )

    indices, values = get_sae_activations(
        sae, 
        get_loader_from_dict([activation_cache], batch_size), 
        device, 
        False, 
        no_verbose,
    )
    return indices, values, toks_cache, activation_cache

def rank_top_act(
    sae: SAE_Template | SAE_LENS, 
    values: Tensor, 
    indices: Tensor,
    k: int,
    activation_threshold: float = 1e-3,
) -> Tensor:
    '''
    Rank the features activate the most frequent, given the values and indices
    '''
    count = t.bincount(indices.flatten(), weights=(values > activation_threshold).flatten())
    return t.topk(count, k).indices
    
def top_act_on_target_components(
    sae: SAE_Template | SAE_LENS,
    model: HookedTransformer,
    toks: Tensor,
    target_components: List[Tuple[Node, Index]],
    k : int,
    batch_size: int,
    device: str,
    activation_threshold: float = 1e-3, 
) -> Tensor:
    '''
    Return the top activation features on the target components
    '''
    indices, values, _, _  = get_cache_with_toks_dataset_on_target_components(
        sae,
        model, 
        toks, 
        target_components,
        batch_size,
        device,
    )

    return rank_top_act(sae, values, indices, k=k, activation_threshold=activation_threshold)

def find_top_activating_components(
    sae: SAE_Template | SAE_LENS,
    activation_cache: ActivationCache | Dict[str, Tensor],
    batch_size: int,
    device: str,
    filter_threshold: float = 1e-3,
    filter_resid: bool = True,
) -> Dict[int, List[Tuple[str, float]]]:
    '''
    Get the top activating components for all feature id.
    Return: Dict[feature_id, List[(name, score)]], the list is sorted by score
    '''
    num_features = sae.cfg.d_sae
    activated_components = {feat_id : {} for feat_id in range(num_features)}
    for name, data in tqdm(activation_cache.items()):
        if "resid" in name and filter_resid:
            continue
        indices_comp, values_comp = get_sae_activations(
            sae, 
            get_loader(data, batch_size), 
            device, 
            False, 
            no_tqdm=True
        )
        for feat_id in range(num_features):
            score = t.where(
                indices_comp == feat_id, (values_comp > filter_threshold), t.zeros_like(values_comp)
            ).sum().item()
            if score > 0:
                activated_components[feat_id][name] = score
    return {
        feat_id: sorted(activated_components[feat_id].items(), key=lambda x: x[1], reverse=True)
        for feat_id in range(num_features)
    }
    
def fetch_max_activating_examples(
    sae: SAE_Template | SAE_LENS,
    model: HookedTransformer,
    feat_idx: int,
    values: Tensor,
    indices: Tensor,
    toks: Tensor,
    k: int = 10,
    buffer: int = 10,
    show_input: bool = True,
) -> list[tuple[float, list[str], list[str], int, int]]:
    """
    Returns the max activating examples across a number of batches from the activations store.
    """
    # filter the activations features
    max_acts_per_row = t.max(t.where(indices == feat_idx, values, t.zeros_like(values)), dim=-1).values
    
    # Get largest indices, get the corresponding max acts, and get the surrounding indices
    k_largest_indices = get_k_largest_indices(max_acts_per_row, k=k, buffer=buffer)
    tokens_with_buffer = index_with_buffer(toks, k_largest_indices, buffer=buffer)
    str_toks = [model.to_str_tokens(tok) for tok in tokens_with_buffer]
    str_input = [model.to_str_tokens(toks[row]) if show_input else "" for row, _ in k_largest_indices]
    top_acts = index_with_buffer(max_acts_per_row, k_largest_indices).tolist()
    data = list(zip(top_acts, str_toks, str_input, [buffer] * len(str_toks), [col for _, col in k_largest_indices]))

    return sorted(data, key=lambda x: x[0], reverse=True)[:k] # type: ignore

from jaxtyping import Float, Int
import html

# --- Helper Functions ---

def get_k_largest_indices(
    x: t.Tensor, 
    k: int, 
    buffer: int = 0, 
    no_overlap: bool = True
) -> t.Tensor:
    """
    Extracts the top k indices from tensor x.
    
    Improvements:
    1. Does NOT filter out edges (start/end of sequence).
    2. Safely handles overlap masking at edges.
    3. Stops early if values are 0 (prevents picking low/zero activation padding).
    """
    # Work on a copy so we can mask out neighbors without affecting original
    temp_x = x.clone()
    
    indices = []
    batch_size, seq_len = x.shape
    
    for _ in range(k):
        # 1. Find the global maximum in the remaining tensor
        # We flatten to use argmax, then convert back to (row, col)
        flat_idx = temp_x.argmax()
        max_val = temp_x.flatten()[flat_idx].item()
        
        # 2. Stop if the maximum value is 0 (or negative)
        # This prevents returning indices that are just empty padding
        if max_val <= 0:
            break
            
        row = flat_idx // seq_len
        col = flat_idx % seq_len
        
        indices.append([row, col])
        
        # 3. Mask out the selected token and its neighbors
        if no_overlap:
            # Dynamic clamping to prevent index out of bounds errors
            # This allows us to pick index 0 even if buffer is 5
            start = max(0, col - buffer)
            end = min(seq_len, col + buffer + 1)
            temp_x[row, start:end] = 0.0
        else:
            temp_x[row, col] = 0.0
            
    if len(indices) == 0:
        return t.empty((0, 2), dtype=t.long)
        
    return t.tensor(indices, dtype=t.long)


def index_with_buffer(
    x: Float[Tensor, "batch seq"], indices: Int[Tensor, "k 2"], buffer: int | None = None
) -> Float[Tensor, "k *buffer_x2_plus1"]:
    if indices.shape[0] == 0:
        return t.tensor([], device=x.device)

    rows, cols = indices.unbind(dim=-1)
    if buffer is not None:
        cols = cols.clone()
        cols = t.clamp(cols, min=buffer, max=x.size(1) - buffer - 1)
        offsets = t.arange(-buffer, buffer + 1, device=x.device)
        rows_expanded = einops.repeat(rows, "k -> k w", w=len(offsets))
        cols_expanded = einops.repeat(cols, "k -> k w", w=len(offsets)) + offsets
        return x[rows_expanded, cols_expanded]
        
    return x[rows, cols]

# --- Updated Data Extraction & Visualization ---

def fetch_feature_activation_intervals(
    model: Any, 
    feat_idx: int,
    values: t.Tensor, 
    indices: t.Tensor, 
    toks: t.Tensor, 
    n_intervals: int = 5,
    k_per_interval: int = 5,
    buffer: int = 5, 
) -> Dict[str, Any]:
    """
    Extracts examples for specific feature activation intervals.
    """
    
    # 1. Reconstruct dense activations
    feat_acts = t.zeros_like(toks, dtype=t.float)
    mask = (indices == feat_idx)
    feat_acts = (values * mask.float()).sum(dim=-1)
    
    max_act = feat_acts.max().item()
    
    # Debug output
    if max_act <= 0:
        return {"feature_id": feat_idx, "intervals": [], "max_act": 0}

    boundaries = t.linspace(1e-3, max_act, n_intervals + 1)
    interval_data = []
    total_tokens = feat_acts.numel()
    seq_len = toks.shape[1]

    for i in range(n_intervals, 0, -1):
        lower_bound = boundaries[i-1].item()
        upper_bound = boundaries[i].item()
        
        # Inclusive upper bound for the top interval
        if i == n_intervals:
            interval_mask = (feat_acts >= lower_bound) & (feat_acts <= upper_bound + 1e-4)
        else:
            interval_mask = (feat_acts >= lower_bound) & (feat_acts < upper_bound)
            
        count = interval_mask.sum().item()
        density = count / total_tokens
        
        if count == 0:
            continue

        # Apply mask to activations so get_k_largest only sees valid interval data
        masked_acts = feat_acts * interval_mask.float()
        
        # Call the FIXED selector
        # We pass the buffer here so it knows how wide to mask neighbors,
        # but the function now safely handles edges.
        selected_indices = get_k_largest_indices(masked_acts, k=k_per_interval, buffer=buffer, no_overlap=True)
        
        if selected_indices.shape[0] == 0:
            continue
        
        examples = []
        for j in range(selected_indices.shape[0]):
            row_idx = selected_indices[j, 0].item()
            col_idx = selected_indices[j, 1].item()
            
            actual_val = feat_acts[row_idx, col_idx].item() # type: ignore

            # --- Dynamic Window Slicing ---
            # Calculate start/end based on buffer, clamped to sequence limits.
            # This removes the need for padding tokens.
            start_idx = max(0, col_idx - buffer)
            end_idx = min(seq_len, col_idx + buffer + 1)
            
            # 1. Short View (Windowed)
            raw_ids_short = toks[row_idx, start_idx:end_idx].tolist() # type: ignore
            str_tokens_short = model.to_str_tokens(t.tensor(raw_ids_short))
            acts_short = feat_acts[row_idx, start_idx:end_idx].tolist() # type: ignore
            
            # Calculate target index relative to the short slice
            # e.g., if slice starts at 5 and target is at 7, relative is 2
            relative_target_idx = col_idx - start_idx
            
            # 2. Full View
            raw_ids_full = toks[row_idx].tolist() # type: ignore
            acts_full = feat_acts[row_idx].tolist() # type: ignore
            str_tokens_full = model.to_str_tokens(t.tensor(raw_ids_full))
            
            examples.append({
                "short": {
                    "tokens": str_tokens_short,
                    "activations": acts_short,
                    "target_idx": relative_target_idx
                },
                "full": {
                    "tokens": str_tokens_full,
                    "activations": acts_full,
                    "target_idx": col_idx
                },
                "target_act": actual_val,
            })

        if len(examples) > 0:
            interval_data.append({
                "min": lower_bound,
                "max": upper_bound,
                "density": density,
                "examples": examples
            })

    return {
        "feature_id": feat_idx,
        "max_act": max_act,
        "intervals": interval_data
    }

def _render_token_html(tokens: List[str], acts: List[float], target_idx: int, max_global: float) -> str:
    """
    Helper to generate the HTML string for a sequence of tokens with green highlighting.
    """
    html_str = ""
    for i, (tok, act) in enumerate(zip(tokens, acts)):
        # 1. Calculate opacity relative to global max for consistent coloring across examples
        # Using a slight non-linear scaling (sqrt) can sometimes help see lower values, 
        # but linear is standard. We'll stick to linear here.
        ratio = act / max_global if max_global > 0 else 0
        alpha = max(0.0, min(1.0, ratio))
        
        # 2. Use RGBA for the Green highlight (Matches the original style)
        # rgba(0, 200, 83) is a vibrant green. 
        # We add a base opacity of 0.1 for non-zero items so they are visible, 
        # or just pure alpha. Let's use pure alpha for accuracy.
        bg_color = f"rgba(0, 200, 83, {alpha})" if act > 0 else "transparent"
        
        # Escape HTML characters (like <, >) and replace newlines with a visual symbol
        safe_tok = html.escape(tok).replace('\n', '↵')
        
        # 3. Target highlighting (The specific token the SAE looked at)
        is_target = (i == target_idx)
        target_class = "target-token" if is_target else ""
        
        # 4. Construct the span with a custom tooltip
        # REMOVE THE NEWLINE in the f-string - it's causing each token to be on a new line
        html_str += f'<span class="token {target_class}" style="background-color: {bg_color};">{safe_tok}<span class="tooltip">{act:.4f}</span></span>'
        
    return html_str

def generate_interactive_html(data: Dict[str, Any]) -> str:
    css = """
    <style>
        body { 
            font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif; 
            background-color: #f8f9fa; 
            color: #333; 
            padding: 20px; 
        }
        .container { 
            max-width: 1100px; 
            margin: 0 auto; 
            background: white; 
            border-radius: 8px; 
            box-shadow: 0 2px 10px rgba(0,0,0,0.05); 
            overflow: hidden; 
        }
        .header { 
            background: #fff; 
            padding: 15px 20px; 
            border-bottom: 1px solid #eee; 
        }
        .header h2 { margin: 0; font-size: 1.1rem; color: #444; }
        
        .interval-block { border-bottom: 1px solid #eee; }
        .interval-header { 
            background: #f1f3f5; 
            padding: 5px 20px; 
            font-size: 0.75rem; 
            font-weight: bold; 
            color: #555; 
            text-transform: uppercase; 
            letter-spacing: 0.5px;
            display: flex;
            justify-content: space-between;
        }
        
        /* Compact Row Layout */
        .example-row { 
            display: flex; 
            align-items: flex-start; 
            padding: 6px 15px; /* Reduced whitespace */
            border-bottom: 1px solid #f9f9f9; 
        }
        .example-row:last-child { border-bottom: none; }
        
        .meta-col { 
            width: 60px; 
            flex-shrink: 0; 
            display: flex; 
            flex-direction: column; 
            align-items: center; 
            margin-right: 15px;
            padding-top: 2px;
        }
        .act-badge { 
            background: #e9ecef; 
            color: #495057; 
            padding: 1px 5px; 
            border-radius: 3px; 
            font-size: 0.65rem; 
            font-weight: bold;
            margin-bottom: 2px;
            text-transform: uppercase;
        }
        .act-val { color: #28a745; font-weight: bold; font-size: 0.85rem; }
        
        .seq-col { 
            flex-grow: 1; 
            font-family: 'Consolas', 'Monaco', monospace; 
            font-size: 13px;
            line-height: 1.1;
            white-space: pre-wrap; 
            cursor: pointer; 
        }
        
        /* Token Styling */
        .token { 
            display: inline; 
            padding: 0 1px;      
            margin: 0;
            border-radius: 2px;
            cursor: help; 
            position: relative; /* Needed for tooltip positioning */
            transition: outline 0.1s;
        }
        
        /* Hover Effect: Outline + Tooltip Visibility */
        .token:hover {
            outline: 1px solid #444;
            z-index: 10;
        }
        
        .target-token { 
            border-bottom: 2px solid #333; 
            font-weight: bold;
        }
        
        /* Custom CSS Tooltip */
        .token .tooltip {
            visibility: hidden;
            background-color: #333;
            color: #fff;
            text-align: center;
            padding: 4px 8px;
            border-radius: 4px;
            position: absolute;
            z-index: 100;
            bottom: 120%; /* Position above the token */
            left: 50%;
            transform: translateX(-50%);
            font-family: sans-serif;
            font-size: 11px;
            font-weight: normal;
            white-space: nowrap;
            pointer-events: none;
            opacity: 0;
            transition: opacity 0.2s;
            box-shadow: 0 2px 5px rgba(0,0,0,0.2);
        }
        
        /* Tiny arrow for the tooltip */
        .token .tooltip::after {
            content: "";
            position: absolute;
            top: 100%;
            left: 50%;
            margin-left: -5px;
            border-width: 5px;
            border-style: solid;
            border-color: #333 transparent transparent transparent;
        }

        .token:hover .tooltip {
            visibility: visible;
            opacity: 1;
        }
        
        .seq-col:hover {
            background-color: #fafafa; /* Subtle hint that row is clickable */
        }

    </style>
    <script>
        function toggleSeq(element) {
            const shortView = element.querySelector('.short-view');
            const fullView = element.querySelector('.full-view');
            
            if (shortView.style.display !== 'none') {
                shortView.style.display = 'none';
                fullView.style.display = 'block';
            } else {
                shortView.style.display = 'block';
                fullView.style.display = 'none';
            }
        }
    </script>
    """
    
    html_content = f"""
    <!DOCTYPE html>
    <html>
    <head>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>SAE Feature {data['feature_id']}</title>
        {css}
    </head>
    <body>
        <div class="container">
            <div class="header">
                <h2>Feature {data['feature_id']} (Max Act: {data['max_act']:.4f})</h2>
            </div>
    """
    
    max_global = data['max_act']
    
    for interval in data['intervals']:
        density_pct = interval['density'] * 100
        
        html_content += f"""
            <div class="interval-block">
                <div class="interval-header">
                    <span>Interval {interval['min']:.3f} - {interval['max']:.3f}</span>
                    <span>Density: {density_pct:.4f}%</span>
                </div>
        """
        
        for ex in interval['examples']:
            short_html = _render_token_html(
                ex['short']['tokens'], 
                ex['short']['activations'], 
                ex['short']['target_idx'], 
                max_global
            )
            
            full_html = _render_token_html(
                ex['full']['tokens'], 
                ex['full']['activations'], 
                ex['full']['target_idx'], 
                max_global
            )
            
            html_content += f"""
                <div class="example-row">
                    <div class="meta-col">
                        <span class="act-badge">Max</span>
                        <span class="act-val">{ex['target_act']:.2f}</span>
                    </div>
                    <div class="seq-col" onclick="toggleSeq(this)">
                        <div class="short-view">{short_html}</div>
                        <div class="full-view" style="display:none;">{full_html}</div>
                    </div>
                </div>
            """
            
        html_content += "</div>"
        
    html_content += """
        </div>
    </body>
    </html>
    """
    
    return html_content

def sparsity(
    feat_idx: int,
    values: Tensor,
    indices: Tensor,
    threshold: float = 1e-3,
) -> float:
    """
    Returns the sparsity of the activations for a particular feature.
    """
    act = t.where((indices == feat_idx) & (values > threshold), t.ones_like(values), t.zeros_like(values))
    batch, seq = indices.shape[:-1]
    return act.sum().item() / (batch * seq)


def extract_data(
    dataloader: DataLoader[tuple[t.Tensor] | Dict[str, t.Tensor]], 
    num_batches: int = 200,
) -> Tensor:
    '''
    Get the tokens from dataset up to given number of batches
    '''
    data = []
    for i, batch in enumerate(dataloader):
        if i >= num_batches:
            break
        if isinstance(batch, dict):
            data.append(batch['input_ids'])
        else:
            data.append(batch[0])
    return t.cat(data, dim=0)

def extract_lvlm_data(
    dataloader: DataLoader,
    num_batches: int = 200,
) -> Dict[str, List[Tensor]]:
    '''
    Get the tokens from lvlm dataset up to given number of batches
    '''
    data = dict()
    for i, batch in enumerate(dataloader):
        if i >= num_batches:
            break
        data['pixel_values']=data.get('pixel_values', []) + [batch['pixel_values']]
        data['input_ids']=data.get('input_ids', []) + [batch['input_ids']]
        data['attention_mask']=data.get('attention_mask', []) + [batch['attention_mask']]
        data['imgids']=data.get('imgids', []) + [batch['imgids']]
        data['captions']=data.get('captions', []) + [batch['captions']]
    
    # for key in data.keys():
    #     print(data[key])
    #     data[key] = t.cat(data[key], dim=0)
    return data


def show_top_logits(
    model: HookedTransformer,
    sae: SAE_Template | SAE_LENS,
    latent_idx: int,
    k: int = 10,
) -> None:
    """
    Displays the top & bottom logits for a particular latent.
    """
    if isinstance(sae, SAE_LENS):
        logits = model.ln_final(sae.W_dec[latent_idx]) @ model.W_U
    else:
        logits = model.ln_final(sae.decoder.weight.data.T[latent_idx]) @ model.W_U

    pos_logits, pos_token_ids = logits.topk(k)
    pos_tokens = model.to_str_tokens(pos_token_ids)
    neg_logits, neg_token_ids = logits.topk(k, largest=False)
    neg_tokens = model.to_str_tokens(neg_token_ids)

    print(
        tabulate(
            zip(map(repr, neg_tokens), neg_logits, map(repr, pos_tokens), pos_logits),
            headers=["Bottom tokens", "Value", "Top tokens", "Value"],
            tablefmt="simple_outline",
            stralign="right",
            numalign="left",
            floatfmt="+.3f",
        )
    )
    
def create_prompt(
    sae: SAE_Template | SAE_LENS,
    model: HookedTransformer,
    latent_idx: int,
    values: Tensor,
    indices: Tensor,
    toks: Tensor,
    k: int = 15,
    buffer: int = 10,
) -> dict[str, str]:
    """
    Returns the system, user & assistant prompts for autointerp.
    """
    data = fetch_max_activating_examples(sae, model, latent_idx, values, indices, toks, k, buffer)
    str_formatted_examples = "\n".join(
        f"{i+1}. {''.join(f'<<{tok}>>' if j == buffer else tok for j, tok in enumerate(seq[1]))}"
        for i, seq in enumerate(data)
    )

    return {
        "system": """We're studying neurons in a neural network. Each neuron activates on some particular word or concept in a short document. The activating words in each document are indicated with << ... >>. Look at the parts of the document the neuron activates for and summarize in a single sentence what the neuron is activating on. Try to be specific in your explanations, although don't be so specific that you exclude some of the examples from matching your explanation. Pay attention to things like the capitalization and punctuation of the activating words or concepts, if that seems relevant. Keep the explanation as short and simple as possible, limited to 20 words or less. Omit punctuation and formatting. You should avoid giving long lists of words.""",
        "user": f"""The activating documents are given below:\n\n{str_formatted_examples}""",
        "assistant": "this neuron fires on",
    }
    

from openai import OpenAI
from dotenv import load_dotenv
import os

def get_autointerp_explanation(
    sae: SAE_Template | SAE_LENS,
    model: HookedTransformer,
    latent_idx: int,
    values: Tensor,
    indices: Tensor,
    toks: Tensor,
    k: int = 15,
    buffer: int = 10,
    n_completions: int = 1,
) -> list[str]:
    """
    Queries OpenAI's API using prompts returned from `create_prompt`, and returns
    a list of the completions.
    """
    client = OpenAI(api_key=API_KEY)

    prompts = create_prompt(sae, model, latent_idx, values, indices, toks, k, buffer)

    result = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": prompts["system"]},
            {"role": "user", "content": prompts["user"]},
            {"role": "assistant", "content": prompts["assistant"]},
        ],
        n=n_completions,
        max_tokens=50,
        stream=False,
    )
    return [choice.message.content for choice in result.choices] # type: ignore

load_dotenv()

API_KEY = os.getenv("OPENAI_API_KEY", None)