model/llava/hooked_llava.py

import json
import logging
from contextlib import contextmanager
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import torch
from jaxtyping import Float
from transformer_lens.ActivationCache import ActivationCache
from transformer_lens.hook_points import HookPoint  # Hooking utilities
from transformer_lens.HookedTransformer import HookedTransformer
from .modeling_llava import *

from sae_lens.sae import SAE

SingleLoss = Float[torch.Tensor, ""]  # Type alias for a single element tensor
LossPerToken = Float[torch.Tensor, "batch pos-1"]
Loss = Union[SingleLoss, LossPerToken]


def get_deep_attr(obj: Any, path: str):
    """Helper function to get a nested attribute from a object.
    In practice used to access HookedTransformer HookPoints (eg model.blocks[0].attn.hook_z)

    Args:
        obj: Any object. In practice, this is a HookedTransformer (or subclass)
        path: str. The path to the attribute you want to access. (eg "blocks.0.attn.hook_z")

    returns:
        Any. The attribute at the end of the path
    """
    parts = path.split(".")
    # Navigate to the last component in the path
    for part in parts:
        obj = obj[int(part)] if part.isdigit() else getattr(obj, part)
    return obj


def set_deep_attr(obj: Any, path: str, value: Any):
    """Helper function to change the value of a nested attribute from a object.
    In practice used to swap HookedTransformer HookPoints (eg model.blocks[0].attn.hook_z) 
    with HookedSAEs and vice versa

    Args:
        obj: Any object. In practice, this is a HookedTransformer (or subclass)
        path: str. The path to the attribute you want to access. (eg "blocks.0.attn.hook_z")
        value: Any. The value you want to set the attribute to (eg a HookedSAE object)
    """
    parts = path.split(".")
    # Navigate to the last component in the path
    for part in parts[:-1]:
        obj = obj[int(part)] if part.isdigit() else getattr(obj, part)
    # Set the value on the final attribute
    setattr(obj, parts[-1], value)

class HookedSAELlavaConditionalGeneration(LlavaForConditionalGeneration):
    def __init__(
        self,
        *model_args: Any,
        **model_kwargs: Any,
    ):
        """Model initialization. Just HookedTransformer init, but adds a dictionary to keep track of attached SAEs.

        Note that if you want to load the model from pretrained weights, you should use
        :meth:`from_pretrained` instead.

        Args:
            *model_args: Positional arguments for HookedTransformer initialization
            **model_kwargs: Keyword arguments for HookedTransformer initialization
        """
        super().__init__(*model_args, **model_kwargs)
        self.acts_to_saes: Dict[str, SAE] = {}

    def add_sae(self, sae: SAE, use_error_term: Optional[bool] = None):
        """Attaches an SAE to the model

        WARNING: This sae will be permanantly attached until you remove it with reset_saes. 
        This function will also overwrite any existing SAE attached to the same hook point.

        Args:
            sae: SparseAutoencoderBase. The SAE to attach to the model
            use_error_term: (Optional[bool]) If provided, will set the use_error_term attribute of the SAE to this value. 
            Determines whether the SAE returns input or reconstruction. Defaults to None.
        """
        act_name = sae.cfg.hook_name
        if (act_name not in self.acts_to_saes) and (act_name not in self.hook_dict):
            logging.warning(
                f"No hook found for {act_name}. Skipping. Check model.hook_dict for available hooks."
            )
            return

        if use_error_term is not None:
            if not hasattr(sae, "_original_use_error_term"):
                sae._original_use_error_term = sae.use_error_term  # type: ignore
            sae.use_error_term = use_error_term
        self.acts_to_saes[act_name] = sae
        set_deep_attr(self, act_name, sae)
        self.setup()

    def _reset_sae(self, act_name: str, prev_sae: Optional[SAE] = None):
        """Resets an SAE that was attached to the model

        By default will remove the SAE from that hook_point.
        If prev_sae is provided, will replace the current SAE with the provided one.
        This is mainly used to restore previously attached SAEs after temporarily running with different SAEs (eg with run_with_saes)

        Args:
            act_name: str. The hook_name of the SAE to reset
            prev_sae: Optional[HookedSAE]. The SAE to replace the current one with. If None, will just remove the SAE from this hook point. Defaults to None
        """
        if act_name not in self.acts_to_saes:
            logging.warning(
                f"No SAE is attached to {act_name}. There's nothing to reset."
            )
            return

        current_sae = self.acts_to_saes[act_name]
        if hasattr(current_sae, "_original_use_error_term"):
            current_sae.use_error_term = current_sae._original_use_error_term
            delattr(current_sae, "_original_use_error_term")

        if prev_sae:
            set_deep_attr(self, act_name, prev_sae)
            self.acts_to_saes[act_name] = prev_sae
        else:
            set_deep_attr(self, act_name, HookPoint())
            del self.acts_to_saes[act_name]

    def reset_saes(
        self,
        act_names: Optional[Union[str, List[str]]] = None,
        prev_saes: Optional[List[Union[SAE, None]]] = None,
    ):
        """Reset the SAEs attached to the model

        If act_names are provided will just reset SAEs attached to those hooks. Otherwise will reset all SAEs attached to the model.
        Optionally can provide a list of prev_saes to reset to. This is mainly used to restore previously attached SAEs after temporarily running with different SAEs (eg with run_with_saes).

        Args:
            act_names (Optional[Union[str, List[str]]): The act_names of the SAEs to reset. If None, will reset all SAEs attached to the model. Defaults to None.
            prev_saes (Optional[List[Union[HookedSAE, None]]]): List of SAEs to replace the current ones with. If None, will just remove the SAEs. Defaults to None.
        """
        if isinstance(act_names, str):
            act_names = [act_names]
        elif act_names is None:
            act_names = list(self.acts_to_saes.keys())

        if prev_saes:
            if len(act_names) != len(prev_saes):
                raise ValueError("act_names and prev_saes must have the same length")
        else:
            prev_saes = [None] * len(act_names)  # type: ignore

        for act_name, prev_sae in zip(act_names, prev_saes):  # type: ignore
            self._reset_sae(act_name, prev_sae)

        self.setup()

    def run_with_saes(
        self,
        *model_args: Any,
        saes: Union[SAE, List[SAE]] = [],
        reset_saes_end: bool = True,
        use_error_term: Optional[bool] = None,
        **model_kwargs: Any,
    ) -> Union[
        None,
        Float[torch.Tensor, "batch pos d_vocab"],
        Loss,
        Tuple[Float[torch.Tensor, "batch pos d_vocab"], Loss],
    ]:
        """Wrapper around HookedTransformer forward pass.

        Runs the model with the given SAEs attached for one forward pass, then removes them. By default, will reset all SAEs to original state after.

        Args:
            *model_args: Positional arguments for the model forward pass
            saes: (Union[HookedSAE, List[HookedSAE]]) The SAEs to be attached for this forward pass
            reset_saes_end (bool): If True, all SAEs added during this run are removed at the end, and previously attached SAEs are restored to their original state. Default is True.
            use_error_term: (Optional[bool]) If provided, will set the use_error_term attribute of all SAEs attached during this run to this value. Defaults to None.
            **model_kwargs: Keyword arguments for the model forward pass
        """
        with self.saes(
            saes=saes, reset_saes_end=reset_saes_end, use_error_term=use_error_term
        ):
            return self(*model_args, **model_kwargs)

    def run_with_cache_with_saes(
        self,
        *model_args: Any,
        saes: Union[SAE, List[SAE]] = [],
        reset_saes_end: bool = True,
        use_error_term: Optional[bool] = None,
        remove_batch_dim: bool = False,
        **kwargs: Any,
    ) -> Tuple[
        Union[
            None,
            Float[torch.Tensor, "batch pos d_vocab"],
            Loss,
            Tuple[Float[torch.Tensor, "batch pos d_vocab"], Loss],
        ],
        Union[ActivationCache, Dict[str, torch.Tensor]],
    ]:
        """Wrapper around 'run_with_cache' in HookedTransformer.

        Attaches given SAEs before running the model with cache and then removes them.
        By default, will reset all SAEs to original state after.

        Args:
            *model_args: Positional arguments for the model forward pass
            saes: (Union[HookedSAE, List[HookedSAE]]) The SAEs to be attached for this forward pass
            reset_saes_end: (bool) If True, all SAEs added during this run are removed at the end, and previously attached SAEs are restored to their original state. Default is True.
            use_error_term: (Optional[bool]) If provided, will set the use_error_term attribute of all SAEs attached during this run to this value. Determines whether the SAE returns input or reconstruction. Defaults to None.
            remove_batch_dim: (bool) Whether to remove the batch dimension (only works for batch_size==1). Defaults to False.
            **kwargs: Keyword arguments for the model forward pass
        """
        with self.saes(
            saes=saes, reset_saes_end=reset_saes_end, use_error_term=use_error_term
        ):
            return self.run_with_cache(  # type: ignore
                *model_args,
                remove_batch_dim=remove_batch_dim,
                **kwargs,
            )

    def run_with_hooks_with_saes(
        self,
        *model_args: Any,
        saes: Union[SAE, List[SAE]] = [],
        reset_saes_end: bool = True,
        fwd_hooks: List[Tuple[Union[str, Callable], Callable]] = [],  # type: ignore
        bwd_hooks: List[Tuple[Union[str, Callable], Callable]] = [],  # type: ignore
        reset_hooks_end: bool = True,
        clear_contexts: bool = False,
        # name_filter: Optional[Callable[[str], bool]] = None,
        **model_kwargs: Any,
    ):
        """Wrapper around 'run_with_hooks' in HookedTransformer.

        Attaches the given SAEs to the model before running the model with hooks and then removes them.
        By default, will reset all SAEs to original state after.

        Args:
            *model_args: Positional arguments for the model forward pass
            act_names: (Union[HookedSAE, List[HookedSAE]]) The SAEs to be attached for this forward pass
            reset_saes_end: (bool) If True, all SAEs added during this run are removed at the end, and previously attached SAEs are restored to their original state. (default: True)
            fwd_hooks: (List[Tuple[Union[str, Callable], Callable]]) List of forward hooks to apply
            bwd_hooks: (List[Tuple[Union[str, Callable], Callable]]) List of backward hooks to apply
            reset_hooks_end: (bool) Whether to reset the hooks at the end of the forward pass (default: True)
            clear_contexts: (bool) Whether to clear the contexts at the end of the forward pass (default: False)
            **model_kwargs: Keyword arguments for the model forward pass
        """
        with self.saes(saes=saes, reset_saes_end=reset_saes_end):
            return self.run_with_hooks(
                *model_args,
                fwd_hooks=fwd_hooks,
                bwd_hooks=bwd_hooks,
                reset_hooks_end=reset_hooks_end,
                clear_contexts=clear_contexts,
                # name_filter=name_filter,
                **model_kwargs,
            )

    @contextmanager
    def saes(
        self,
        saes: Union[SAE, List[SAE]] = [],
        reset_saes_end: bool = True,
        use_error_term: Optional[bool] = None,
    ):
        """
        A context manager for adding temporary SAEs to the model.
        See HookedTransformer.hooks for a similar context manager for hooks.
        By default will keep track of previously attached SAEs, and restore them when the context manager exits.

        Example:

        .. code-block:: python

            from transformer_lens import HookedSAETransformer, HookedSAE, HookedSAEConfig

            model = HookedSAETransformer.from_pretrained('gpt2-small')
            sae_cfg = HookedSAEConfig(...)
            sae = HookedSAE(sae_cfg)
            with model.saes(saes=[sae]):
                spliced_logits = model(text)


        Args:
            saes (Union[HookedSAE, List[HookedSAE]]): SAEs to be attached.
            reset_saes_end (bool): If True, removes all SAEs added by this context manager when the context manager exits, returning previously attached SAEs to their original state.
            use_error_term (Optional[bool]): If provided, will set the use_error_term attribute of all SAEs attached during this run to this value. Defaults to None.
        """
        act_names_to_reset = []
        prev_saes = []
        if isinstance(saes, SAE):
            saes = [saes]
        try:
            for sae in saes:
                act_names_to_reset.append(sae.cfg.hook_name)
                prev_sae = self.acts_to_saes.get(sae.cfg.hook_name, None)
                prev_saes.append(prev_sae)
                self.add_sae(sae, use_error_term=use_error_term)
            yield self
        finally:
            if reset_saes_end:
                self.reset_saes(act_names_to_reset, prev_saes)
                        
    def get_activation_at(self, inputs, hook_point):
        act_cache = {}

        def hook_fn(act, hook):
            act_cache[hook.name] = act.detach().cpu()
            raise EarlyExit()

        try:
            self.run_with_hooks(inputs, fwd_hooks=[(hook_point, hook_fn)])
        except EarlyExit:
            pass

        return act_cache[hook_point]
    
class EarlyExit(Exception):
    pass


_NULLU_DOWN_PROJ_PREFIXES = (
    "model.layers",                      # haotian-liu LLaVA-v1.5 (Nullu default)
    "language_model.model.layers",       # older llava-hf transformers format
    "language_model.layers",             # newer llava-hf transformers format
)
_NULLU_DOWN_PROJ_SUFFIX = ".mlp.down_proj.weight"


def _match_down_proj_layer(key: str) -> Optional[int]:
    """Return the layer index if `key` is an `mlp.down_proj.weight` under a known prefix."""
    if not key.endswith(_NULLU_DOWN_PROJ_SUFFIX):
        return None
    for prefix in _NULLU_DOWN_PROJ_PREFIXES:
        if key.startswith(prefix + "."):
            middle = key[len(prefix) + 1 : -len(_NULLU_DOWN_PROJ_SUFFIX)]
            if middle.isdigit():
                return int(middle)
    return None


def _load_edited_down_proj_weights(
    edited_model_path: Union[str, Path], layer_indices: List[int]
) -> Dict[int, torch.Tensor]:
    """Read only `mlp.down_proj.weight` tensors for `layer_indices` from a saved HF checkpoint.

    Supports sharded / single-file safetensors and pytorch_model.bin variants. Uses
    `safe_open` so only the requested tensors are materialized.
    """
    from safetensors.torch import safe_open

    edited_path = Path(edited_model_path)
    wanted = set(layer_indices)
    result: Dict[int, torch.Tensor] = {}

    shard_files = sorted(edited_path.glob("*.safetensors"))
    if shard_files:
        for shard in shard_files:
            with safe_open(shard, framework="pt") as f:
                for key in f.keys():
                    idx = _match_down_proj_layer(key)
                    if idx is not None and idx in wanted and idx not in result:
                        result[idx] = f.get_tensor(key)
            if len(result) == len(wanted):
                break
    else:
        bin_files = sorted(edited_path.glob("pytorch_model*.bin")) or sorted(
            edited_path.glob("*.bin")
        )
        if not bin_files:
            raise FileNotFoundError(
                f"No safetensors or .bin weight files found in {edited_path}"
            )
        for bf in bin_files:
            sd = torch.load(bf, map_location="cpu", weights_only=True)
            for key, tensor in sd.items():
                idx = _match_down_proj_layer(key)
                if idx is not None and idx in wanted and idx not in result:
                    result[idx] = tensor
            if len(result) == len(wanted):
                break

    missing = wanted - set(result.keys())
    if missing:
        raise KeyError(
            f"Could not find mlp.down_proj.weight for layers {sorted(missing)} in {edited_path}"
        )
    return result


def load_nullu_model(
    lowest_layer: int,
    highest_layer: int,
    edited_model_path: Union[str, Path],
    base_model_name: str = "llava-hf/llava-1.5-7b-hf",
    torch_dtype: torch.dtype = torch.float16,
    device: Union[str, torch.device] = "cuda:0",
) -> "HookedSAELlavaConditionalGeneration":
    """Load a `HookedSAELlavaConditionalGeneration` with Nullu edits applied only to a layer slice.

    Nullu's `scripts/model_edit.py` saves a checkpoint where every layer in `[0, n_layers)` has had
    its `mlp.down_proj.weight` null-projected onto the non-hallucination subspace. This helper loads
    that full-range edit and transplants only the `[lowest_layer, highest_layer)` slice onto an
    otherwise-unedited LLaVA model — avoiding a fresh edit run for each range you want to study.

    Args:
        lowest_layer: inclusive lower bound of edited layer range.
        highest_layer: exclusive upper bound of edited layer range.
        edited_model_path: Nullu's saved edited model directory, e.g.
            ``Nullu/output/edited_model/LLaVA-7B-top4-0-32-test``.
        base_model_name: unedited LLaVA HF model id or local path.
        torch_dtype: dtype for model weights.
        device: device to load the model on.

    Returns:
        A `HookedSAELlavaConditionalGeneration` with edited down_proj weights on the chosen layers.
    """
    model = HookedSAELlavaConditionalGeneration.from_pretrained(
        base_model_name, torch_dtype=torch_dtype
    ).to(device)

    n_layers = model.config.text_config.num_hidden_layers
    if not (0 <= lowest_layer < highest_layer <= n_layers):
        raise ValueError(
            f"Require 0 <= lowest_layer < highest_layer <= {n_layers}; "
            f"got lowest_layer={lowest_layer}, highest_layer={highest_layer}"
        )

    layer_indices = list(range(lowest_layer, highest_layer))
    edited_weights = _load_edited_down_proj_weights(edited_model_path, layer_indices)

    language_layers = model.model.language_model.layers
    for idx, w in edited_weights.items():
        tgt = language_layers[idx].mlp.down_proj.weight
        tgt.data.copy_(w.to(device=tgt.device, dtype=tgt.dtype))

    model.eval()
    return model