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 : Zero hardcoding , is_visualizable_module any module (zero hard coding), labeling or legend , x-axis , y-axis , names properly title , properly 


import os
from typing import Any, Dict, Union, Optional
import logging
from pathlib import Path

import torch
import torch.nn as nn
import plotly.graph_objs as go
from plotly.subplots import make_subplots
import numpy as np

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class ModelVisualizer:
    """
    A class for visualizing tensors from various model layers.
    """

    def __init__(self, output_dir: Union[str, Path] = "visualizations"):
        """
        Initialize the ModelVisualizer.

        Args:
            output_dir (Union[str, Path]): Directory to save visualization outputs.
        """
        self.output_dir = Path(output_dir)
        self.html_dir = self.output_dir / "html"
        self.png_dir = self.output_dir / "png"
        self.html_dir.mkdir(parents=True, exist_ok=True)
        self.png_dir.mkdir(parents=True, exist_ok=True)
        logger.info(f"Visualizations will be saved in: {self.output_dir}")

    def visualize_tensor(self, tensor: torch.Tensor, metadata: Dict[str, Any]) -> None:
        """
        Visualize a tensor based on its dimensionality.

        Args:
            tensor (torch.Tensor): The tensor to visualize.
            metadata (Dict[str, Any]): Additional information for visualization.
        """
        dim = tensor.dim()
        logger.info(f"Visualizing a tensor with {dim} dimensions.")

        try:
            if dim == 1:
                self._plot_1d(tensor, metadata)
            elif dim == 2:
                self._plot_2d(tensor, metadata)
            elif dim == 3:
                self._plot_3d(tensor, metadata)
            elif dim == 4:
                self._plot_4d(tensor, metadata)
            else:
                logger.warning(f"Tensors with {dim} dimensions are not supported for visualization.")
        except Exception as e:
            logger.error(f"Error in visualizing tensor: {e}", exc_info=True)
            raise

    def _plot_1d(self, tensor: torch.Tensor, metadata: Dict[str, Any]) -> None:
        """
        Generate a 1D scatter plot.

        Args:
            tensor (torch.Tensor): 1D tensor to plot.
            metadata (Dict[str, Any]): Metadata for the plot title and labels.
        """
        x = np.arange(tensor.size(0))
        y = tensor.cpu().numpy()

        fig = go.Figure(data=go.Scatter(x=x, y=y, mode='markers', name='1D Scatter'))
        title = self._generate_title(metadata, "1D Scatter Plot")
        fig.update_layout(
            title=title,
            xaxis_title='Index',
            yaxis_title='Value',
            template='plotly_white'
        )
        plot_name = self._sanitize_filename(f"{metadata.get('layer_name', 'layer')}_1d_scatter")
        self._save_and_export(fig, plot_name)

    def _plot_2d(self, tensor: torch.Tensor, metadata: Dict[str, Any]) -> None:
        """
        Generate a 2D scatter plot.

        Args:
            tensor (torch.Tensor): 2D tensor to plot.
            metadata (Dict[str, Any]): Metadata for the plot title and labels.
        """
        if tensor.size(0) < 2:
            logger.warning("2D tensor does not have at least two dimensions for scatter plot.")
            return

        x = tensor[0].cpu().numpy()
        y = tensor[1].cpu().numpy()

        fig = go.Figure(data=go.Scatter(x=x, y=y, mode='markers', name='2D Scatter'))
        title = self._generate_title(metadata, "2D Scatter Plot")
        fig.update_layout(
            title=title,
            xaxis_title='Dimension 1',
            yaxis_title='Dimension 2',
            template='plotly_white'
        )
        plot_name = self._sanitize_filename(f"{metadata.get('layer_name', 'layer')}_2d_scatter")
        self._save_and_export(fig, plot_name)

    def _plot_3d(self, tensor: torch.Tensor, metadata: Dict[str, Any]) -> None:
        """
        Generate a 3D scatter plot.

        Args:
            tensor (torch.Tensor): 3D tensor to plot.
            metadata (Dict[str, Any]): Metadata for the plot title and labels.
        """
        if tensor.size(0) < 3:
            logger.warning("3D tensor does not have at least three dimensions for scatter plot.")
            return

        x = tensor[0].cpu().numpy()
        y = tensor[1].cpu().numpy()
        z = tensor[2].cpu().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(metadata, "3D Scatter Plot")
        fig.update_layout(
            title=title,
            scene=dict(
                xaxis_title='Dimension 1',
                yaxis_title='Dimension 2',
                zaxis_title='Dimension 3'
            ),
            template='plotly_white'
        )
        plot_name = self._sanitize_filename(f"{metadata.get('layer_name', 'layer')}_3d_scatter")
        self._save_and_export(fig, plot_name)

    def _plot_4d(self, tensor: torch.Tensor, metadata: Dict[str, Any]) -> None:
        """
        Generate a 4D scatter plot by incorporating color as the fourth dimension.

        Args:
            tensor (torch.Tensor): 4D tensor to plot.
            metadata (Dict[str, Any]): Metadata for the plot title and labels.
        """
        if tensor.size(0) < 3:
            logger.warning("4D tensor does not have at least three dimensions for scatter plot.")
            return

        x = tensor[0].cpu().numpy()
        y = tensor[1].cpu().numpy()
        z = tensor[2].cpu().numpy()
        color = tensor[3].cpu().numpy()

        fig = go.Figure(data=[go.Scatter3d(
            x=x, y=y, z=z,
            mode='markers',
            marker=dict(
                size=5,
                color=color,
                colorscale='Viridis',
                colorbar=dict(title='Dimension 4'),
                showscale=True
            ),
            name='4D Scatter'
        )])
        title = self._generate_title(metadata, "4D Scatter Plot")
        fig.update_layout(
            title=title,
            scene=dict(
                xaxis_title='Dimension 1',
                yaxis_title='Dimension 2',
                zaxis_title='Dimension 3'
            ),
            template='plotly_white'
        )
        plot_name = self._sanitize_filename(f"{metadata.get('layer_name', 'layer')}_4d_scatter")
        self._save_and_export(fig, plot_name)

    def _generate_title(self, metadata: Dict[str, Any], plot_type: str) -> str:
        """
        Generate a descriptive title based on metadata.

        Args:
            metadata (Dict[str, Any]): Metadata containing model and layer information.
            plot_type (str): Type of the plot.

        Returns:
            str: Generated title.
        """
        model_type = metadata.get('model_type', 'Unknown Model')
        layer_name = metadata.get('layer_name', 'Unknown Layer')
        return f"{plot_type} - {model_type}, Layer: {layer_name}"

    def _sanitize_filename(self, name: str) -> str:
        """
        Sanitize the filename by replacing spaces and illegal characters.

        Args:
            name (str): Original filename.

        Returns:
            str: Sanitized filename.
        """
        return "".join(c if c.isalnum() or c in ('_', '-') else '_' for c in name)

    def _save_and_export(self, fig: go.Figure, plot_name: str) -> None:
        """
        Save the figure as HTML and PNG, and log the actions.

        Args:
            fig (go.Figure): Plotly figure to save.
            plot_name (str): Base name for the saved files.
        """
        try:
            html_path = self.html_dir / f"{plot_name}.html"
            png_path = self.png_dir / f"{plot_name}.png"
            fig.write_html(str(html_path))
            fig.write_image(str(png_path))
            logger.info(f"Plot saved as HTML: {html_path}")
            logger.info(f"Plot saved as PNG: {png_path}")
        except Exception as e:
            logger.error(f"Failed to save plot {plot_name}: {e}", exc_info=True)
            raise


class AdvancedModelWrapper(nn.Module):
    """
    A wrapper class for advanced model visualization.
    """

    def __init__(self, model: nn.Module, visualizer: ModelVisualizer):
        """
        Initialize the AdvancedModelWrapper.

        Args:
            model (nn.Module): The PyTorch model to wrap.
            visualizer (ModelVisualizer): The visualizer instance.
        """
        super().__init__()
        self.model = model
        self.visualizer = visualizer
        self.hooks = []
        self._register_hooks()

    def _register_hooks(self) -> None:
        """
        Register forward hooks for all suitable modules in the model.
        """
        for name, module in self.model.named_modules():
            if self._is_visualizable_module(module):
                hook = module.register_forward_hook(self._make_hook(name))
                self.hooks.append(hook)
                logger.debug(f"Hook registered for layer: {name}")

    @staticmethod
    def _is_visualizable_module(module: nn.Module) -> bool:
        """
        Determine if the module is suitable for visualization.

        Args:
            module (nn.Module): The module to check.

        Returns:
            bool: True if visualizable, False otherwise.
        """
        return isinstance(module, (nn.Conv2d, nn.Linear, nn.LSTM, nn.Transformer, nn.ReLU, nn.BatchNorm2d))

    def _make_hook(self, name: str):
        """
        Create a hook function for the given module name.

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

        Returns:
            Callable: The hook function.
        """
        def hook(module: nn.Module, input: Any, output: Any) -> None:
            metadata = {
                'model_type': type(self.model).__name__,
                'layer_name': name,
                'layer_type': type(module).__name__
            }
            if isinstance(output, torch.Tensor):
                self.visualizer.visualize_tensor(output, metadata)
            elif isinstance(output, (list, tuple)):
                for idx, tensor in enumerate(output):
                    if isinstance(tensor, torch.Tensor):
                        meta = metadata.copy()
                        meta['output_index'] = idx
                        self.visualizer.visualize_tensor(tensor, meta)
            else:
                logger.warning(f"Unsupported output type from layer {name}: {type(output)}")

        return hook

    def forward(self, *args, **kwargs) -> Any:
        """
        Forward pass through the wrapped model.

        Args:
            *args: Variable length argument list.
            **kwargs: Arbitrary keyword arguments.

        Returns:
            Any: The output of the wrapped model.
        """
        try:
            return self.model(*args, **kwargs)
        except Exception as e:
            logger.error(f"Error during forward pass: {e}", exc_info=True)
            raise

    def remove_hooks(self) -> None:
        """
        Remove all registered hooks.
        """
        for hook in self.hooks:
            hook.remove()
        logger.info("All hooks have been removed.")
        self.hooks.clear()


def visualize_model(model: nn.Module, input_data: torch.Tensor, output_dir: str = "visualizations") -> None:
    """
    Visualize the internals of a PyTorch model.

    Args:
        model (nn.Module): The PyTorch model to visualize.
        input_data (torch.Tensor): Sample input data for the model.
        output_dir (str): Directory to save visualization outputs.
    """
    visualizer = ModelVisualizer(output_dir)
    wrapped_model = AdvancedModelWrapper(model, visualizer)

    try:
        logger.info("Starting model visualization...")
        wrapped_model.eval()
        with torch.no_grad():
            wrapped_model(input_data)
        logger.info("Model visualization completed successfully.")
    except Exception as e:
        logger.error(f"Error during model visualization: {e}", exc_info=True)
    finally:
        wrapped_model.remove_hooks()




def main():
    """
    Main function to demonstrate model visualization.
    """
    from transformers import AutoModelForCausalLM
    model=AutoModelForCausalLM.from_pretrained("gpt2")
 
    # Create sample input data
    input_tensor = torch.randint(50257,(10, 512) ,dtype=torch.long)  

    # Visualize the model
    visualize_model(model, input_tensor)


if __name__ == "__main__":
    main()