Total number of parameters: 0
Loaded model 'gpt2' with 0 parameters.
Model has no parameters. Exiting.

import sys
from typing import Any, Dict, List, Optional, Tuple

import torch
from torch.nn import Module
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    PreTrainedModel,
    PreTrainedTokenizer,
)
import plotly.graph_objs as go
from plotly.subplots import make_subplots


def load_model_and_tokenizer(
    model_name: str,
) -> Tuple[PreTrainedModel, PreTrainedTokenizer]:
    """
    Load the pre-trained model and tokenizer.

    Args:
        model_name (str): Name of the pre-trained model.

    Returns:
        Tuple containing the model and tokenizer.
    """
    try:
        model = AutoModelForCausalLM.from_pretrained(model_name)
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model.eval()
        return model, tokenizer
    except Exception as e:
        print(f"Error loading model or tokenizer: {e}", file=sys.stderr)
        raise


def get_neuron_parameters(
    model: PreTrainedModel, recurse: bool = False
) -> Dict[str, Dict[str, Any]]:
    """
    Retrieve model parameters with their metadata.

    Args:
        model (PreTrainedModel): The pre-trained model.
        recurse (bool): Whether to recurse through submodules.

    Returns:
        Dictionary mapping parameter names to their metadata.
    """
    parameters = {}
    try:
        param_count = 0
        for name, param in model.named_parameters(recurse=recurse):
            shape = tuple(param.shape)
            trainable = param.requires_grad
            num_params = param.numel()
            parameters[name] = {
                "shape": shape,
                "trainable": trainable,
                "num_params": num_params,
            }
            param_count += num_params
        print(f"Total number of parameters: {param_count}")
        return parameters
    except Exception as e:
        print(f"Error retrieving parameters: {e}", file=sys.stderr)
        raise


class ActivationCapturer:
    """
    Captures activations from specified layers using forward hooks.
    """

    def __init__(self, model: PreTrainedModel, layers: Optional[List[str]] = None):
        """
        Initialize the capturer.

        Args:
            model (PreTrainedModel): The pre-trained model.
            layers (List[str], optional): List of layer names to capture.
                If None, captures all transformer layers.
        """
        self.model = model
        self.layers = layers or [
            name for name, _ in model.named_modules() if "layer" in name
        ]
        self.activations: Dict[str, torch.Tensor] = {}
        self.hooks: List[torch.utils.hooks.RemovableHandle] = []
        self._register_hooks()

    def _register_hooks(self) -> None:
        """
        Register forward hooks to capture activations.
        """
        try:
            for name, module in self.model.named_modules():
                if name in self.layers:
                    hook = module.register_forward_hook(self._get_hook(name))
                    self.hooks.append(hook)
        except Exception as e:
            print(f"Error registering hooks: {e}", file=sys.stderr)
            raise

    def _get_hook(self, name: str):
        """
        Create a hook function to capture activations.

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

        Returns:
            Hook function.
        """

        def hook(module: Module, input: Tuple[Any, ...], output: torch.Tensor):
            self.activations[name] = output.detach()

        return hook

    def clear(self) -> None:
        """
        Clear captured activations.
        """
        self.activations.clear()

    def remove_hooks(self) -> None:
        """
        Remove all registered hooks.
        """
        for hook in self.hooks:
            hook.remove()


def compute_neuron_attributions(
    model: PreTrainedModel,
    tokenizer: PreTrainedTokenizer,
    input_text: str,
    target_token_index: Optional[int] = None,
    top_k: int = 10,
) -> Optional[Dict[str, torch.Tensor]]:
    """
    Compute neuron attributions using gradients.

    Args:
        model (PreTrainedModel): The pre-trained model.
        tokenizer (PreTrainedTokenizer): The tokenizer.
        input_text (str): The input string.
        target_token_index (int, optional): The index of the target token
            to analyze. If None, analyzes the last token.
        top_k (int): Number of top neurons to identify per layer.

    Returns:
        Dictionary mapping layer names to their top influencing neurons.
    """
    try:
        inputs = tokenizer(input_text, return_tensors="pt")
        input_ids = inputs["input_ids"]
        attention_mask = inputs.get("attention_mask", None)

        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model.to(device)
        input_ids = input_ids.to(device)
        if attention_mask is not None:
            attention_mask = attention_mask.to(device)

        # Initialize activation capturer
        capturer = ActivationCapturer(model)
        model.zero_grad()

        # Ensure gradients are enabled
        for param in model.parameters():
            param.requires_grad = True

        with torch.enable_grad():
            outputs = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                output_hidden_states=True,
                return_dict=True,
            )
            logits = outputs.logits
        # print(logits)

        # Determine target token index
        if target_token_index is None:
            target_token_index = inputs["input_ids"].shape[1] - 1

        # Select logits for the target token
        target_logits = logits[0, target_token_index]
        target_token = torch.argmax(target_logits)

        # Compute loss as the logit of the target token
        loss = target_logits[target_token]
        loss.backward()

        # Compute attributions
        attributions = {}
        for layer_name, activation in capturer.activations.items():
            if activation.requires_grad:
                gradient = activation.grad
                if gradient is None:
                    print(
                        f"No gradient found for layer '{layer_name}'. Skipping attribution.",
                        file=sys.stderr,
                    )
                    continue
                # Compute the importance of each neuron
                importance = torch.abs(gradient).mean(dim=(0, 2, 3))  # Adjust dims as needed
                top_importance, top_indices = torch.topk(importance, top_k)
                attributions[layer_name] = top_importance.cpu()

        capturer.remove_hooks()
        return attributions

    except Exception as e:
        print(f"Error computing neuron attributions: {e}", file=sys.stderr)
        return None


def visualize_attributions(
    attributions: Dict[str, torch.Tensor], input_text: str, top_k: int
) -> None:
    """
    Visualize neuron attributions using Plotly.

    Args:
        attributions (Dict[str, torch.Tensor]): Neuron attributions per layer.
        input_text (str): The input string.
        top_k (int): Number of top neurons visualized per layer.
    """
    try:
        layers = list(attributions.keys())
        num_layers = len(layers)
        fig = make_subplots(
            rows=num_layers,
            cols=1,
            subplot_titles=layers,
            vertical_spacing=0.02,
        )

        for idx, layer in enumerate(layers, start=1):
            data = attributions[layer].numpy()
            neurons = list(range(len(data)))
            fig.add_trace(
                go.Bar(
                    x=neurons,
                    y=data,
                    marker=dict(color=data, colorscale="Viridis"),
                    showlegend=False,
                ),
                row=idx,
                col=1,
            )

            fig.update_yaxes(title_text="Attribution", row=idx, col=1)
            fig.update_xaxes(title_text="Neuron Index", row=idx, col=1)

        fig.update_layout(
            title=f"Top {top_k} Neuron Attributions for Input: \"{input_text}\"",
            height=400 * num_layers,
            showlegend=False,
            hovermode="closest",
        )
        fig.show()
    except Exception as e:
        print(f"Error during visualization: {e}", file=sys.stderr)
        raise


def main(input_text: str, model_name: str = "gpt2", top_k: int = 10) -> None:
    """
    Main function to execute the workflow.

    Args:
        input_text (str): The input string to process.
        model_name (str): The name of the pre-trained model.
        top_k (int): Number of top neurons to identify per layer.
    """
    try:
        model, tokenizer = load_model_and_tokenizer(model_name)
        params = get_neuron_parameters(model)
        print(
            f"Loaded model '{model_name}' with {sum(p['num_params'] for p in params.values())} parameters.",
            file=sys.stdout,
        )

        if not params:
            print("Model has no parameters. Exiting.", file=sys.stderr)
            return

        attributions = compute_neuron_attributions(
            model, tokenizer, input_text, top_k=top_k
        )
        if attributions is None:
            print("Failed to compute neuron attributions.", file=sys.stderr)
            return

        visualize_attributions(attributions, input_text, top_k)
    except Exception as e:
        print(f"An error occurred in the main workflow: {e}", file=sys.stderr)


if __name__ == "__main__":
    # Example usage:
    example_input = "Once upon a time"
    main(example_input)