main.py

import argparse
import gc
import json
import os
import random
import sys
import time
from dataclasses import asdict, dataclass
from datetime import datetime
from typing import Any, Protocol

import torch
import torch.nn.functional as F
from collectibles import ListCollection
from einops import rearrange
from loguru import logger
from sae_lens import SAE, ActivationsStore
from sae_lens.sae import TopK
from torch import nn
from tqdm import tqdm
from transformer_lens import HookedTransformer

from sae_bench.evals.mdl.eval_config import MDLEvalConfig
from sae_bench.sae_bench_utils import (
    activation_collection,
    general_utils,
    get_eval_uuid,
    get_sae_bench_version,
    get_sae_lens_version,
)
from sae_bench.sae_bench_utils.sae_selection_utils import (
    get_saes_from_regex,
)

EVAL_TYPE = "mdl"


class Decodable(Protocol):
    def decode(self, x: torch.Tensor) -> torch.Tensor: ...


def build_bins(
    min_pos_activations_F: torch.Tensor,
    max_activations_F: torch.Tensor,
    bin_precision: float | None = None,  # 0.2,
    num_bins: int | None = None,  # 16)
) -> list[torch.Tensor]:
    if bin_precision is not None and num_bins is not None:
        raise ValueError("Only one of bin_precision or num_bins should be provided")
    if bin_precision is None and num_bins is None:
        raise ValueError("Either bin_precision or num_bins should be provided")

    num_features = len(max_activations_F)

    assert len(max_activations_F) == num_features

    # positive_mask_BsF = feature_activations_BsF > 0
    # masked_activations_BsF = torch.where(positive_mask_BsF, feature_activations_BsF, torch.inf)
    # min_pos_activations_F = torch.min(masked_activations_BsF, dim=-1).values
    # min_pos_activations_F = torch.where(
    #     torch.isfinite(min_pos_activations_F), min_pos_activations_F, 0
    # )
    min_pos_activations_F = torch.zeros_like(max_activations_F)

    logger.debug(max_activations_F)
    logger.debug(min_pos_activations_F)

    bins_F_list_Bi: list[torch.Tensor] = []

    if bin_precision is not None:
        for feature_idx in range(num_features):
            bins = torch.arange(
                min_pos_activations_F[feature_idx].item(),
                max_activations_F[feature_idx].item() + 2 * bin_precision,
                bin_precision,
                device=max_activations_F.device,
            )
            bins_F_list_Bi.append(bins)

        return bins_F_list_Bi

    else:
        assert num_bins is not None
        for feature_idx in range(num_features):
            bins = torch.linspace(
                min_pos_activations_F[feature_idx].item(),
                max_activations_F[feature_idx].item(),
                num_bins + 1,
                device=max_activations_F.device,
            )
            bins_F_list_Bi.append(bins)

        return bins_F_list_Bi


def calculate_dl(
    num_features: int,
    bins_F_list_Bi: list[torch.Tensor],
    device: str,
    activations_store: ActivationsStore,
    sae: SAE,
    k: int,
) -> float:
    float_entropy_F = torch.zeros(num_features, device=device, dtype=torch.float32)
    bool_entropy_F = torch.zeros(num_features, device=device, dtype=torch.float32)

    x_BSN = activations_store.get_buffer(config.sae_batch_size)[0]
    feature_activations_BsF = sae.encode(x_BSN).squeeze()

    if feature_activations_BsF.ndim == 2:
        feature_activations_BsF = feature_activations_BsF
    elif feature_activations_BsF.ndim == 3:
        feature_activations_BsF = rearrange(
            feature_activations_BsF,
            "batch seq_len num_features -> (batch seq_len) num_features",
        )
    else:
        raise ValueError("feature_activations should be 2D or 3D tensor")

    for feature_idx in tqdm(range(num_features), desc="Calculating DL"):
        # BOOL entropy
        bool_prob = torch.zeros(1, device=device)

        bool_prob_F = (feature_activations_BsF > 0).float().mean(dim=0)
        bool_prob = bool_prob + bool_prob_F[feature_idx]

        if bool_prob == 0 or bool_prob == 1:
            bool_entropy = 0
        else:
            bool_entropy = -bool_prob * torch.log2(bool_prob) - (
                1 - bool_prob
            ) * torch.log2(1 - bool_prob)
        bool_entropy_F[feature_idx] = bool_entropy

        # FLOAT entropy
        num_bins = len(bins_F_list_Bi[feature_idx]) - 1
        counts_Bi = torch.zeros(num_bins, device="cpu")

        feature_activations_Bs = feature_activations_BsF[:, feature_idx].to(
            dtype=torch.float32
        )
        bins = bins_F_list_Bi[feature_idx]

        temp_counts_Bi, _bin_edges = torch.histogram(
            feature_activations_Bs.cpu(), bins=bins.cpu()
        )
        counts_Bi = counts_Bi + temp_counts_Bi

        counts_Bi = counts_Bi.to(device)

        probs_Bi = counts_Bi / counts_Bi.sum()
        probs_Bi = probs_Bi[(probs_Bi > 0) & (probs_Bi < 1)]

        if len(probs_Bi) == 0:
            float_entropy = 0
        else:
            # H[p] = -sum(p * log2(p))
            float_entropy = -torch.sum(probs_Bi * torch.log2(probs_Bi)).item()

        float_entropy_F[feature_idx] = float_entropy

    total_entropy_F = (
        bool_entropy_F.cuda() + bool_prob_F.cuda() * float_entropy_F.cuda()  # type: ignore
    )

    description_length = total_entropy_F.sum().item()

    return description_length


def quantize_features_to_bin_midpoints(
    features_BF: torch.Tensor, bins_F_list_Bi: list[torch.Tensor]
) -> torch.Tensor:
    """
    Quantize features to the bin midpoints of their corresponding histograms.
    """
    _, num_features = features_BF.shape

    quantized_features_BF = torch.empty_like(features_BF, device=features_BF.device)

    for feature_idx in range(num_features):
        # Extract the feature values and bin edges for the current histogram
        features_B = features_BF[:, feature_idx]
        bin_edges_Bi = bins_F_list_Bi[feature_idx]

        num_bins = len(bin_edges_Bi) - 1

        bin_indices_B = torch.bucketize(features_B, bin_edges_Bi)
        bin_indices_clipped_B = torch.clamp(bin_indices_B, min=1, max=num_bins) - 1

        # Calculate the midpoints of the bins
        bin_mids_Bi = 0.5 * (bin_edges_Bi[:-1] + bin_edges_Bi[1:])

        quantized_features_BF[:, feature_idx] = bin_mids_Bi[bin_indices_clipped_B]

    return quantized_features_BF


# def calculate_dl(
#     activations_store: ActivationsStore,
#     sae: SAE,
#     bins: list[torch.Tensor],
#     k: int | None = None,
# ) -> float:
#     for i in range(10):
#         x_BSN = activations_store.get_buffer(config.sae_batch_size)
#         feature_activations_BsF = sae.encode(x_BSN).squeeze()

#         if feature_activations_BsF.ndim == 2:
#             feature_activations_BsF = feature_activations_BsF
#         elif feature_activations_BsF.ndim == 3:
#             feature_activations_BsF = rearrange(
#                 feature_activations_BsF,
#                 "batch seq_len num_features -> (batch seq_len) num_features",
#             )
#         else:
#             raise ValueError("feature_activations should be 2D or 3D tensor")

#         if k is not None:
#             topk_fn = TopK(k)
#             feature_activations_BsF = topk_fn(feature_activations_BsF)

#         entropy = _calculate_dl_single(feature_activations_BsF, bins)
#     return entropy


def check_quantised_features_reach_mse_threshold(
    bins_F_list_Bi: list[torch.Tensor],
    activations_store: ActivationsStore,
    sae: SAE,
    mse_threshold: float,
    autoencoder: SAE,
    k: int | None = None,
) -> tuple[bool, float]:
    mse_losses: list[torch.Tensor] = []

    for i in range(1):
        x_BSN = activations_store.get_buffer(config.sae_batch_size)[0]
        feature_activations_BSF = sae.encode(x_BSN).squeeze()

        if k is not None:
            topk_fn = TopK(k)
            feature_activations_BSF = topk_fn(feature_activations_BSF)

        quantised_feature_activations_BsF = quantize_features_to_bin_midpoints(
            feature_activations_BSF, bins_F_list_Bi
        )

        reconstructed_x_BSN: torch.Tensor = autoencoder.decode(
            quantised_feature_activations_BsF
        )

        mse_loss: torch.Tensor = F.mse_loss(
            reconstructed_x_BSN, x_BSN.squeeze(), reduction="mean"
        )
        mse_loss = torch.sqrt(mse_loss) / sae.cfg.d_in
        mse_losses.append(mse_loss)

    avg_mse_loss = torch.mean(torch.stack(mse_losses))
    within_threshold = bool((avg_mse_loss < mse_threshold).item())

    return within_threshold, mse_loss.item()  # type: ignore


class IdentityAE(nn.Module):
    def forward(self, x):
        return x

    def decode(self, x):
        return x


@dataclass
class MDLEvalResult:
    num_bins: int
    bins: list[torch.Tensor]
    k: int | None

    description_length: float
    within_threshold: bool
    mse_loss: float

    def to_dict(self) -> dict[str, Any]:
        out = asdict(self)
        out["bins"] = []
        return out


class MDLEvalResultsCollection(ListCollection[MDLEvalResult]):
    num_bins: list[int]
    bins: list[list[torch.Tensor]]
    k: list[int] | None

    description_length: list[float]
    within_threshold: list[bool]
    mse_loss: list[float]

    def pick_minimum_viable(self) -> MDLEvalResult:
        all_description_lengths = torch.tensor(self.description_length)
        threshold_mask = torch.tensor(self.within_threshold)

        viable_description_lengths = all_description_lengths[threshold_mask]
        if len(viable_description_lengths) > 0:
            min_dl_idx = int(torch.argmin(viable_description_lengths).item())
            return self[min_dl_idx]

        else:
            min_dl_idx = int(torch.argmin(all_description_lengths).item())
            return self[min_dl_idx]


def run_eval_single_sae(
    config: MDLEvalConfig,
    sae: SAE,
    model: HookedTransformer,
    device: str,
    dataset_name: str = "HuggingFaceFW/fineweb",
) -> MDLEvalResultsCollection:
    random.seed(config.random_seed)
    torch.manual_seed(config.random_seed)

    torch.set_grad_enabled(False)
    mdl_eval_results_list: list[MDLEvalResult] = []

    sae.cfg.dataset_trust_remote_code = True
    sae = sae.to(device)
    model = model.to(device)  # type: ignore

    activations_store = ActivationsStore.from_sae(
        model, sae, config.sae_batch_size, dataset=dataset_name, device=device
    )

    num_features = sae.cfg.d_sae

    def get_min_max_activations() -> tuple[torch.Tensor, torch.Tensor]:
        min_pos_activations_1F = torch.zeros(1, num_features, device=device)
        max_activations_1F = torch.zeros(1, num_features, device=device) + 100

        for _ in range(10):
            neuron_activations_BSN = activations_store.get_buffer(
                config.sae_batch_size
            )[0]

            feature_activations_BsF = sae.encode(neuron_activations_BSN).squeeze()

            cat_feature_activations_BsF = torch.cat(
                [
                    feature_activations_BsF,
                    min_pos_activations_1F,
                    max_activations_1F,
                ],
                dim=0,
            )
            min_pos_activations_1F = torch.min(
                cat_feature_activations_BsF, dim=0
            ).values.unsqueeze(0)
            max_activations_1F = torch.max(
                cat_feature_activations_BsF, dim=0
            ).values.unsqueeze(0)

        min_pos_activations_F = min_pos_activations_1F.squeeze()
        max_activations_F = max_activations_1F.squeeze()

        return min_pos_activations_F, max_activations_F

    min_pos_activations_F, max_activations_F = get_min_max_activations()

    print("num_bins_values", config.num_bins_values)
    print("k_values", config.k_values)

    for num_bins in config.num_bins_values:
        for k in config.k_values:  # type: ignore
            bins = build_bins(
                min_pos_activations_F, max_activations_F, num_bins=num_bins
            )

            print("Built bins")

            within_threshold, mse_loss = check_quantised_features_reach_mse_threshold(
                bins_F_list_Bi=bins,
                activations_store=activations_store,
                sae=sae,
                mse_threshold=config.mse_epsilon_threshold,
                autoencoder=sae,
                k=k,
            )
            if not within_threshold:
                logger.warning(
                    f"mse_loss for num_bins = {num_bins} and k = {k} is {mse_loss}, which is not within threshold"
                )

            print("Checked threshold")

            description_length = calculate_dl(
                num_features=num_features,
                bins_F_list_Bi=bins,
                device=device,
                activations_store=activations_store,
                sae=sae,
                k=k,
            )

            logger.info(
                f"Description length: {description_length} for num_bins = {num_bins} and k = {k} and mse = {mse_loss}"
            )

            mdl_eval_results_list.append(
                MDLEvalResult(
                    num_bins=num_bins,
                    bins=bins,
                    k=k,
                    description_length=description_length,
                    within_threshold=within_threshold,
                    mse_loss=mse_loss,
                )
            )

    mdl_eval_results = MDLEvalResultsCollection(mdl_eval_results_list)

    result = []

    for mdl_eval_result in mdl_eval_results:
        result.append(mdl_eval_result.to_dict())

    return result  # type: ignore

    # minimum_viable_eval_result = mdl_eval_results.pick_minimum_viable()

    # minimum_viable_description_length = minimum_viable_eval_result.description_length
    # logger.info(minimum_viable_description_length)

    # return minimum_viable_eval_result


def run_eval(
    config: MDLEvalConfig,
    selected_saes: list[tuple[str, SAE]] | list[tuple[str, str]],
    device: str,
    output_path: str,
    force_rerun: bool = False,
) -> dict[str, Any]:
    """
    selected_saes is a list of either tuples of (sae_lens release, sae_lens id) or (sae_name, SAE object)
    """
    eval_instance_id = get_eval_uuid()
    sae_lens_version = get_sae_lens_version()
    sae_bench_commit_hash = get_sae_bench_version()

    results_dict = {}

    llm_dtype = general_utils.str_to_dtype(config.llm_dtype)

    print(f"Using dtype: {llm_dtype}")

    model = HookedTransformer.from_pretrained_no_processing(
        config.model_name, device=device, dtype=llm_dtype
    )

    for sae_release, sae_object_or_id in tqdm(
        selected_saes, desc="Running SAE evaluation on all selected SAEs"
    ):
        sae_id, sae, sparsity = general_utils.load_and_format_sae(
            sae_release, sae_object_or_id, device
        )  # type: ignore
        sae = sae.to(device=device, dtype=llm_dtype)

        sae_result_path = general_utils.get_results_filepath(
            output_path, sae_release, sae_id
        )

        if os.path.exists(sae_result_path) and not force_rerun:
            print(f"Skipping {sae_release}_{sae_id} as results already exist")
            continue

        eval_output = run_eval_single_sae(
            config=config,
            sae=sae,
            model=model,
            dataset_name=config.dataset_name,
            device=device,
        )

        sae_eval_result = {
            "eval_instance_id": eval_instance_id,
            "sae_lens_release": sae_release,
            "sae_lens_id": sae_id,
            "eval_type_id": EVAL_TYPE,
            "sae_lens_version": sae_lens_version,
            "sae_bench_version": sae_bench_commit_hash,
            "date_time": datetime.now().isoformat(),
            "eval_config": asdict(config),
            "eval_results": eval_output,
            "eval_artifacts": {"artifacts": "None"},
            "sae_cfg_dict": asdict(sae.cfg),
        }

        with open(sae_result_path, "w") as f:
            json.dump(sae_eval_result, f, indent=4)

        results_dict[f"{sae_release}_{sae_id}"] = eval_output

    results_dict["custom_eval_config"] = asdict(config)

    gc.collect()
    torch.cuda.empty_cache()

    return results_dict


def create_config_and_selected_saes(
    args,
) -> tuple[MDLEvalConfig, list[tuple[str, str]]]:
    config = MDLEvalConfig(
        model_name=args.model_name,
    )

    if args.llm_batch_size is not None:
        config.llm_batch_size = args.llm_batch_size  # type: ignore
    else:
        config.llm_batch_size = activation_collection.LLM_NAME_TO_BATCH_SIZE[  # type: ignore
            config.model_name
        ]

    if args.llm_dtype is not None:
        config.llm_dtype = args.llm_dtype
    else:
        config.llm_dtype = activation_collection.LLM_NAME_TO_DTYPE[config.model_name]

    if args.random_seed is not None:
        config.random_seed = args.random_seed

    selected_saes = get_saes_from_regex(args.sae_regex_pattern, args.sae_block_pattern)
    assert len(selected_saes) > 0, "No SAEs selected"

    releases = set([release for release, _ in selected_saes])

    print(f"Selected SAEs from releases: {releases}")

    for release, sae in selected_saes:
        print(f"Sample SAEs: {release}, {sae}")

    return config, selected_saes


def arg_parser():
    parser = argparse.ArgumentParser(description="Run MDL evaluation")
    parser.add_argument("--random_seed", type=int, default=None, help="Random seed")
    parser.add_argument("--model_name", type=str, required=True, help="Model name")
    parser.add_argument(
        "--sae_regex_pattern",
        type=str,
        required=True,
        help="Regex pattern for SAE selection",
    )
    parser.add_argument(
        "--sae_block_pattern",
        type=str,
        required=True,
        help="Regex pattern for SAE block selection",
    )
    parser.add_argument(
        "--output_folder",
        type=str,
        default="eval_results/mdl",
        help="Output folder",
    )
    parser.add_argument(
        "--force_rerun", action="store_true", help="Force rerun of experiments"
    )
    parser.add_argument(
        "--clean_up_activations",
        action="store_false",
        help="Clean up activations after evaluation",
    )
    parser.add_argument(
        "--llm_batch_size",
        type=int,
        default=None,
        help="Batch size for LLM. If None, will be populated using LLM_NAME_TO_BATCH_SIZE",
    )
    parser.add_argument(
        "--llm_dtype",
        type=str,
        default=None,
        choices=[None, "float32", "float64", "float16", "bfloat16"],
        help="Data type for LLM. If None, will be populated using LLM_NAME_TO_DTYPE",
    )

    return parser


if __name__ == "__main__":
    """python evals/mdl/main.py \
    --sae_regex_pattern "sae_bench_pythia70m_sweep_standard_ctx128_0712" \
    --sae_block_pattern "blocks.4.hook_resid_post__trainer_10" \
    --model_name pythia-70m-deduped """
    logger.remove()
    logger.add(sys.stdout, level="INFO")

    args = arg_parser().parse_args()
    device = general_utils.setup_environment()

    start_time = time.time()

    config, selected_saes = create_config_and_selected_saes(args)

    print(selected_saes)

    # create output folder
    os.makedirs(args.output_folder, exist_ok=True)

    config = MDLEvalConfig(
        k_values=[None],  # type: ignore
        # num_bins_values=[8, 12, 16, 32, 64, 128],
        num_bins_values=[8, 16, 32, 64],
        # num_bins_values=[8],
        mse_epsilon_threshold=0.2,
        model_name=args.model_name,
    )
    logger.info(config)

    results_dict = run_eval(
        config,
        selected_saes,
        device,
        args.output_folder,
        args.force_rerun,
    )

    end_time = time.time()

    print(f"Finished evaluation in {end_time - start_time} seconds")