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 : below plots are working but coming black, data is missing ,  


import logging
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Union

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

# 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") -> None:
        """
        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)
            elif dim == 5:
                self._plot_5d(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.
        """
        try:
            x = np.arange(tensor.size(0))
            y = tensor.cpu().numpy()

            fig = go.Figure()
            fig.add_trace(go.Scatter(
                x=x,
                y=y,
                mode='markers+lines',
                name=f"{metadata.get('layer_name', 'layer')}_1D_Scatter"
            ))
            title = self._generate_title(metadata, "1D Scatter Plot")
            fig.update_layout(
                title=title,
                xaxis_title='Index',
                yaxis_title='Value',
                template='plotly_white',
                showlegend=True
            )
            plot_name = self._sanitize_filename(f"{metadata.get('layer_name', 'layer')}_1d_scatter")
            self._save_and_export(fig, plot_name)
        except Exception as e:
            logger.error(f"Failed to plot 1D tensor: {e}", exc_info=True)
            raise

    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.
        """
        try:
            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()
            fig.add_trace(go.Scatter(
                x=x,
                y=y,
                mode='markers',
                name=f"{metadata.get('layer_name', 'layer')}_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',
                showlegend=True
            )
            plot_name = self._sanitize_filename(f"{metadata.get('layer_name', 'layer')}_2d_scatter")
            self._save_and_export(fig, plot_name)
        except Exception as e:
            logger.error(f"Failed to plot 2D tensor: {e}", exc_info=True)
            raise

    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.
        """
        try:
            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()
            fig.add_trace(go.Scatter3d(
                x=x,
                y=y,
                z=z,
                mode='markers',
                marker=dict(size=5),
                name=f"{metadata.get('layer_name', 'layer')}_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',
                showlegend=True
            )
            plot_name = self._sanitize_filename(f"{metadata.get('layer_name', 'layer')}_3d_scatter")
            self._save_and_export(fig, plot_name)
        except Exception as e:
            logger.error(f"Failed to plot 3D tensor: {e}", exc_info=True)
            raise

    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.
        """
        try:
            if tensor.size(0) < 4:
                logger.warning("4D tensor does not have at least four 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()
            scatter = 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=f"{metadata.get('layer_name', 'layer')}_4D_Scatter"
            )
            fig.add_trace(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',
                showlegend=True
            )
            plot_name = self._sanitize_filename(f"{metadata.get('layer_name', 'layer')}_4d_scatter")
            self._save_and_export(fig, plot_name)
        except Exception as e:
            logger.error(f"Failed to plot 4D tensor: {e}", exc_info=True)
            raise

    def _plot_5d(self, tensor: torch.Tensor, metadata: Dict[str, Any]) -> None:
        """
        Generate a 5D scatter plot by incorporating color and size as the fourth and fifth dimensions.

        Args:
            tensor (torch.Tensor): 5D tensor to plot.
            metadata (Dict[str, Any]): Metadata for the plot title and labels.
        """
        try:
            if tensor.size(0) < 5:
                logger.warning("5D tensor does not have at least five 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()
            size = tensor[4].cpu().numpy()

            # Normalize size for better visualization
            size_normalized = (size - np.min(size)) / (np.max(size) - np.min(size) + 1e-8) * 20 + 5

            fig = go.Figure()
            scatter = go.Scatter3d(
                x=x,
                y=y,
                z=z,
                mode='markers',
                marker=dict(
                    size=size_normalized,
                    color=color,
                    colorscale='Viridis',
                    colorbar=dict(title='Dimension 4'),
                    showscale=True,
                    sizemode='diameter',
                    opacity=0.8
                ),
                name=f"{metadata.get('layer_name', 'layer')}_5D_Scatter"
            )
            fig.add_trace(scatter)
            title = self._generate_title(metadata, "5D Scatter Plot")
            fig.update_layout(
                title=title,
                scene=dict(
                    xaxis_title='Dimension 1',
                    yaxis_title='Dimension 2',
                    zaxis_title='Dimension 3'
                ),
                template='plotly_white',
                showlegend=True
            )
            # Add size as an annotation
            fig.update_layout(
                annotations=[
                    dict(
                        x=0,
                        y=1.05,
                        showarrow=False,
                        text=f"Dimension 5 (Size): Min={np.min(size):.2f}, Max={np.max(size):.2f}",
                        xref="paper",
                        yref="paper"
                    )
                ]
            )
            plot_name = self._sanitize_filename(f"{metadata.get('layer_name', 'layer')}_5d_scatter")
            self._save_and_export(fig, plot_name)
        except Exception as e:
            logger.error(f"Failed to plot 5D tensor: {e}", exc_info=True)
            raise

    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')
        layer_type = metadata.get('layer_type', 'Unknown Layer Type')
        return f"{plot_type} - Model: {model_type}, Layer: {layer_name} ({layer_type})"

    @staticmethod
    def _sanitize_filename(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) -> None:
        """
        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: List[Callable] = []
        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.
        """
        # Dynamically determine if a module is visualizable based on output tensor dimensions
        # For example, visualize modules that output tensors with no more than 5 dimensions
        # This avoids hardcoding specific module types

        # To-do: Adjust criteria as per specific requirements
        return True  # Allow all modules to be visualized; adjust as needed

    def _make_hook(self, name: str) -> Callable:
        """
        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: Dict[str, Any] = {
                '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: Any, **kwargs: Any) -> 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() -> None:
    """
    Main function to demonstrate model visualization.
    """
    try:
        from transformers import AutoModelForCausalLM
    except ImportError as e:
        logger.error(f"Failed to import transformers: {e}", exc_info=True)
        return

    try:
        model = AutoModelForCausalLM.from_pretrained("gpt2")
        logger.info("Model loaded successfully.")
    except Exception as e:
        logger.error(f"Failed to load model: {e}", exc_info=True)
        return

    # Create sample input data
    # Ensure that the input data matches the model's expected input dimensions and type
    try:
        input_tensor = torch.randint(
            low=0,
            high=50257,
            size=(10, 512),
            dtype=torch.long
        )
        logger.info("Input tensor created successfully.")
    except Exception as e:
        logger.error(f"Failed to create input tensor: {e}", exc_info=True)
        return

    # Visualize the model
    visualize_model(model, input_tensor)


if __name__ == "__main__":
    main()