Following this instruction :
1. Strictly adhere to PEP-8 standards, ensuring code is clean, well-structured, and follows Python's style guide for readability.
2. Use appropriate modules like `typing` for type hints, ensuring code clarity, maintainability, and leveraging Python's standard libraries effectively.
3. Focus on time, space, robustness, scalability; write optimized, efficient code that can handle large datasets and scale seamlessly.
4. Implement comprehensive error handling, catching exceptions gracefully, providing meaningful error messages, and ensuring code robustness under unexpected conditions.
5. Ensure the developer agent writes advanced, maintainable Python code, balancing performance and readability while adhering to best practices.


TASK : use we have advanced develop model visualization tools save in directory(all image all .html format) each and every 1d,2d,3d,4d, 

generalised any model , nlp or computer vision 
properly legend , labelling , colouring plots , very professional ,resuable code

reference code : 

import torch
import torch.nn as nn
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import plotly.graph_objs as go
from plotly.subplots import make_subplots
from typing import Any, Dict
import logging

# Configure logging for debugging and error tracking
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class DynamicTensorModel(nn.Module):
    """
    A PyTorch model that accepts a recursive key input and returns a tensor of varying dimensions.
    It also visualizes the tensor using Plotly based on its dimensionality.
    """

    def __init__(self) -> None:
        super(DynamicTensorModel, self).__init__()

    def forward(self, key_input: Dict[str, Any]) -> torch.Tensor:
        """
        Processes the recursive key input to generate a tensor.

        Args:
            key_input (Dict[str, Any]): A nested dictionary input.

        Returns:
            torch.Tensor: A tensor with dimensions based on the key input.
        """
        try:
            tensor = self.generate_tensor(key_input)
            self.visualize_tensor(tensor, key_input)
            return tensor
        except Exception as e:
            logger.error(f"Error in forward pass: {e}")
            raise

    def generate_tensor(self, key_input: Dict[str, Any], depth: int = 1) -> torch.Tensor:
        """
        Recursively processes the key input to generate a tensor with varying dimensions.

        Args:
            key_input (Dict[str, Any]): A nested dictionary input.
            depth (int, optional): Current depth of recursion. Defaults to 1.

        Returns:
            torch.Tensor: A tensor with dimensions based on recursion depth.
        """
        if not isinstance(key_input, dict):
            raise ValueError("Key input must be a dictionary.")

        if not key_input:
            raise ValueError("Key input dictionary is empty.")

        # Base case: when the value is not a dictionary
        for key, value in key_input.items():
            if isinstance(value, dict):
                return self.generate_tensor(value, depth + 1)
            else:
                # Create a tensor with dimension based on depth
                size = 10  # arbitrary size for each dimension
                dims = min(depth, 4)  # Limit to 4D
                shape = tuple([size] * dims)
                logger.info(f"Generating a {dims}D tensor with shape {shape} based on key '{key}'.")
                return torch.randn(shape)

        # If all values are dictionaries, increase depth
        size = 10
        dims = min(depth, 4)
        shape = tuple([size] * dims)
        logger.info(f"Generating a {dims}D tensor with shape {shape} based on nested keys.")
        return torch.randn(shape)

    def visualize_tensor(self, tensor: torch.Tensor, key_input: Dict[str, Any]) -> None:
        """
        Visualizes the tensor using Plotly based on its dimensionality.

        Args:
            tensor (torch.Tensor): The tensor to visualize.
            key_input (Dict[str, Any]): The original key input for dynamic naming.
        """
        dim = tensor.dim()
        logger.info(f"Visualizing a tensor with {dim} dimensions.")

        if dim == 1:
            self._plot_1d(tensor, key_input)
        elif dim == 2:
            self._plot_2d(tensor, key_input)
        elif dim == 3:
            self._plot_3d(tensor, key_input)
        elif dim == 4:
            self._plot_4d(tensor, key_input)
        else:
            logger.warning(f"Tensors with {dim} dimensions are not supported for visualization.")

    def _plot_1d(self, tensor: torch.Tensor, key_input: Dict[str, Any]) -> None:
        """Generates a simple 1D scatter plot."""
        try:
            x = torch.arange(tensor.size(0)).numpy()
            y = tensor.numpy()
            fig = go.Figure(data=go.Scatter(x=x, y=y, mode='markers', name='1D Scatter'))
            title = self._generate_title(key_input, "1D Scatter Plot")
            fig.update_layout(title=title, xaxis_title='Index', yaxis_title='Value')
            fig.show()
            logger.info("1D scatter plot generated successfully.")
        except Exception as e:
            logger.error(f"Error plotting 1D tensor: {e}")

    def _plot_2d(self, tensor: torch.Tensor, key_input: Dict[str, Any]) -> None:
        """Generates a standard 2D scatter plot."""
        try:
            if tensor.size(1) < 2:
                raise ValueError("2D tensor must have at least two features for plotting.")

            x = tensor[:, 0].numpy()
            y = tensor[:, 1].numpy()
            fig = go.Figure(data=go.Scatter(x=x, y=y, mode='markers', name='2D Scatter'))
            title = self._generate_title(key_input, "2D Scatter Plot")
            fig.update_layout(title=title, xaxis_title='Feature 1', yaxis_title='Feature 2')
            fig.show()
            logger.info("2D scatter plot generated successfully.")
        except Exception as e:
            logger.error(f"Error plotting 2D tensor: {e}")

    def _plot_3d(self, tensor: torch.Tensor, key_input: Dict[str, Any]) -> None:
        """Generates an interactive 3D scatter plot."""
        try:
            if tensor.size(1) < 3:
                raise ValueError("3D tensor must have at least three features for plotting.")

            x = tensor[:, 0].numpy()
            y = tensor[:, 1].numpy()
            z = tensor[:, 2].numpy()
            fig = go.Figure(data=go.Scatter3d(
                x=x, y=y, z=z,
                mode='markers',
                marker=dict(size=5),
                name='3D Scatter'
            ))
            title = self._generate_title(key_input, "3D Scatter Plot")
            fig.update_layout(title=title,
                              scene=dict(
                                  xaxis_title='Feature 1',
                                  yaxis_title='Feature 2',
                                  zaxis_title='Feature 3'
                              ))
            fig.show()
            logger.info("3D scatter plot generated successfully.")
        except Exception as e:
            logger.error(f"Error plotting 3D tensor: {e}")

    def _plot_4d(self, tensor: torch.Tensor, key_input: Dict[str, Any]) -> None:
        """Generates an advanced 3D scatter plot with a dynamic legend based on the key input."""
        try:
            if tensor.size(1) < 4:
                raise ValueError("4D tensor must have at least four features for plotting.")

            x = tensor[:, 0].numpy()
            y = tensor[:, 1].numpy()
            z = tensor[:, 2].numpy()
            color_feature = tensor[:, 3].numpy()

            # Create a color scale based on the fourth dimension
            colorscale = 'Viridis'

            fig = go.Figure(data=go.Scatter3d(
                x=x, y=y, z=z,
                mode='markers',
                marker=dict(
                    size=5,
                    color=color_feature,
                    colorscale=colorscale,
                    opacity=0.8,
                    colorbar=dict(title='Feature 4')
                ),
                name='4D Scatter'
            ))
            title = self._generate_title(key_input, "4D Scatter Plot with Dynamic Legend")
            fig.update_layout(title=title,
                              scene=dict(
                                  xaxis_title='Feature 1',
                                  yaxis_title='Feature 2',
                                  zaxis_title='Feature 3'
                              ))
            fig.show()
            logger.info("4D scatter plot generated successfully.")
        except Exception as e:
            logger.error(f"Error plotting 4D tensor: {e}")

    def _generate_title(self, key_input: Dict[str, Any], plot_type: str) -> str:
        """
        Generates a descriptive title based on the key input.

        Args:
            key_input (Dict[str, Any]): The original key input.
            plot_type (str): The type of plot.

        Returns:
            str: A descriptive title.
        """
        flat_keys = self._flatten_keys(key_input)
        key_str = "_".join(flat_keys)
        title = f"{plot_type} - Keys: {key_str}"
        return title

    def _flatten_keys(self, d: Dict[str, Any], parent_key: str = '', sep: str = '_') -> list:
        """
        Flattens nested dictionary keys.

        Args:
            d (Dict[str, Any]): The dictionary to flatten.
            parent_key (str, optional): The base key string. Defaults to ''.
            sep (str, optional): Separator between keys. Defaults to '_'.

        Returns:
            list: A list of flattened key strings.
        """
        items = []
        for k, v in d.items():
            new_key = f"{parent_key}{sep}{k}" if parent_key else k
            if isinstance(v, dict):
                items.extend(self._flatten_keys(v, new_key, sep=sep))
            else:
                items.append(new_key)
        return items


class GPT2LayerVisualizer:
    """
    A class to visualize activations of each layer in GPT2LMHeadModel using DynamicTensorModel.
    """

    def __init__(self, model: GPT2LMHeadModel, tokenizer: GPT2Tokenizer) -> None:
        self.model = model
        self.tokenizer = tokenizer
        self.dynamic_visualizer = DynamicTensorModel()
        self.activations = {}
        self.register_hooks()

    def register_hooks(self) -> None:
        """
        Registers forward hooks to capture activations from each layer.
        """
        for name, module in self.model.named_modules():
            if isinstance(module, (nn.Linear, nn.Embedding, nn.LayerNorm)):
                module.register_forward_hook(self.get_activation(name))
                logger.info(f"Hook registered for layer: {name}")

    def get_activation(self, name: str):
        """
        Returns a hook function that captures the output of a layer.

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

        Returns:
            function: A hook function.
        """
        def hook(module, input, output):
            self.activations[name] = output.detach().cpu()
            logger.info(f"Activation captured for layer: {name}")
        return hook

    def visualize_activations(self, input_text: str) -> None:
        """
        Processes the input text through the model and visualizes activations.

        Args:
            input_text (str): The input text to process.
        """
        try:
            inputs = self.tokenizer(input_text, return_tensors='pt')
            logger.info(f"Input tokens: {inputs}")

            with torch.no_grad():
                self.model(**inputs)

            for layer_name, activation in self.activations.items():
                key_input = {'layer': layer_name}
                tensor = self.process_activation(activation)
                self.dynamic_visualizer.visualize_tensor(tensor, key_input)

        except Exception as e:
            logger.error(f"Error during visualization: {e}")

    def process_activation(self, activation: torch.Tensor) -> torch.Tensor:
        """
        Processes the activation tensor to match the visualization requirements.

        Since GPT-2 activations are typically high-dimensional (e.g., [batch, sequence, hidden]),
        we'll perform some aggregation or slicing to reduce dimensionality.

        Args:
            activation (torch.Tensor): The activation tensor.

        Returns:
            torch.Tensor: A processed tensor suitable for visualization.
        """
        try:
            # Example processing: take the first element in batch and sequence
            tensor = activation[0, 0, :]

            # If tensor is still multi-dimensional, apply mean across certain dimensions
            if tensor.dim() > 1:
                tensor = tensor.mean(dim=0)

            # Ensure tensor is 1D for simplicity; adjust as needed
            if tensor.dim() > 1:
                tensor = tensor.view(-1)

            return tensor

        except Exception as e:
            logger.error(f"Error processing activation tensor: {e}")
            # Return a random tensor to avoid breaking visualization
            return torch.randn(10)


def main():
    # Initialize tokenizer and model
    tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
    model = GPT2LMHeadModel.from_pretrained('gpt2')
    model.eval()

    # Initialize visualizer
    visualizer = GPT2LayerVisualizer(model, tokenizer)

    # Sample input text
    input_text = "Once upon a time, in a land far away,"

    # Visualize activations
    visualizer.visualize_activations(input_text)


if __name__ == "__main__":
    main()