/usr/local/lib/python3.10/dist-packages/transformers/tokenization_utils_base.py:1601: FutureWarning:

`clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884

ERROR:root:Error during model forward pass: 'tuple' object has no attribute 'detach'
ERROR:root:An error occurred in the main workflow: 'tuple' object has no attribute 'detach'
import logging
from typing import Any, Dict, List, Tuple

import torch
from torch import nn
from transformers import AutoModelForCausalLM, AutoTokenizer
import plotly.graph_objects as go


class NeuronActivationVisualizer:
    """
    A class to compute and visualize neuron activations and their contributions
    to the next token prediction in GPT-2.
    """

    def __init__(self, model_name: str = "gpt2") -> None:
        """
        Initializes the model and tokenizer.

        Args:
            model_name (str): The name of the pre-trained model to load.
        """
        try:
            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
            self.model = AutoModelForCausalLM.from_pretrained(model_name)
            self.model.eval()
            self.activations: Dict[str, torch.Tensor] = {}
            logging.info(f"Model '{model_name}' and tokenizer loaded successfully.")
        except Exception as e:
            logging.error(f"Error loading model or tokenizer: {e}")
            raise

    def _get_activation(self, name: str) -> Any:
        """
        Returns a hook function to capture activations.

        Args:
            name (str): The name of the layer.

        Returns:
            Callable: The hook function.
        """

        def hook(model: nn.Module, input: Tuple[torch.Tensor], output: torch.Tensor) -> None:
            self.activations[name] = output.detach()

        return hook

    def register_hooks(self) -> None:
        """
        Registers hooks to all transformer layers to capture activations.
        """
        try:
            for idx, layer in enumerate(self.model.transformer.h):
                layer_name = f"layer_{idx}"
                layer.register_forward_hook(self._get_activation(layer_name))
            logging.info("Hooks registered successfully.")
        except Exception as e:
            logging.error(f"Error registering hooks: {e}")
            raise

    def get_activations(self, input_ids: torch.Tensor) -> Dict[str, torch.Tensor]:
        """
        Performs a forward pass to obtain activations.

        Args:
            input_ids (torch.Tensor): Tokenized input IDs.

        Returns:
            Dict[str, torch.Tensor]: A dictionary of activations per layer.
        """
        try:
            with torch.no_grad():
                self.model(input_ids)
            return self.activations
        except Exception as e:
            logging.error(f"Error during model forward pass: {e}")
            raise

    def analyze_contributions(self, activations: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
        """
        Analyzes neuron contributions to the next token prediction.

        Args:
            activations (Dict[str, torch.Tensor]): Activations per layer.

        Returns:
            Dict[str, torch.Tensor]: Contributions per layer.
        """
        contributions: Dict[str, torch.Tensor] = {}
        try:
            # Assuming we're interested in the last token's activation
            for layer_name, activation in activations.items():
                # activation shape: [batch_size, sequence_length, hidden_size]
                last_token_activation = activation[:, -1, :]  # Shape: [batch_size, hidden_size]
                # Contribution can be defined as the activation value itself for simplicity
                contributions[layer_name] = last_token_activation.squeeze(0)
            logging.info("Neuron contributions analyzed successfully.")
            return contributions
        except Exception as e:
            logging.error(f"Error analyzing contributions: {e}")
            raise

    def visualize_contributions(self, contributions: Dict[str, torch.Tensor], layers_to_display: int = 12) -> None:
        """
        Visualizes neuron contributions across layers using Plotly.

        Args:
            contributions (Dict[str, torch.Tensor]): Contributions per layer.
            layers_to_display (int): Number of layers to display in the visualization.
        """
        try:
            layer_names = sorted(contributions.keys(), key=lambda x: int(x.split("_")[1]))
            selected_layers = layer_names[:layers_to_display]
            contribution_values = [contributions[layer].cpu().numpy() for layer in selected_layers]

            fig = go.Figure(
                data=go.Heatmap(
                    z=contribution_values,
                    x=list(range(contributions[selected_layers[0]].shape[0])),
                    y=selected_layers,
                    colorscale='Viridis'
                )
            )

            fig.update_layout(
                title="Neuron Contributions per Layer",
                xaxis_title="Neuron Index",
                yaxis_title="Transformer Layers",
                height=600,
                width=1000
            )

            fig.show()
            logging.info("Visualization generated successfully.")
        except Exception as e:
            logging.error(f"Error during visualization: {e}")
            raise

    def process_text(self, text: str) -> torch.Tensor:
        """
        Tokenizes and encodes the input text.

        Args:
            text (str): The input text.

        Returns:
            torch.Tensor: Tokenized input IDs.
        """
        try:
            encoding = self.tokenizer.encode_plus(text, return_tensors='pt')
            input_ids = encoding['input_ids']
            logging.info("Input text tokenized successfully.")
            return input_ids
        except Exception as e:
            logging.error(f"Error processing text: {e}")
            raise

    def main(self, text: str) -> None:
        """
        Orchestrates the workflow from processing input text to visualization.

        Args:
            text (str): The input text.
        """
        try:
            self.register_hooks()
            input_ids = self.process_text(text)
            activations = self.get_activations(input_ids)
            contributions = self.analyze_contributions(activations)
            self.visualize_contributions(contributions)
        except Exception as e:
            logging.error(f"An error occurred in the main workflow: {e}")


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    input_text = (
        "In the realm of artificial intelligence, understanding neuron activations "
        "is crucial for interpreting model decisions and improving performance."
    )
    visualizer = NeuronActivationVisualizer(model_name="gpt2")
    visualizer.main(input_text)