torch-rechub 0.0.4__py3-none-any.whl → 0.0.6__py3-none-any.whl

torch_rechub/utils/model_utils.py

@@ -0,0 +1,233 @@
+"""Common model utility functions for Torch-RecHub.
+
+This module provides shared utilities for model introspection and input generation,
+used by both ONNX export and visualization features.
+
+Examples
+--------
+>>> from torch_rechub.utils.model_utils import extract_feature_info, generate_dummy_input
+>>> feature_info = extract_feature_info(model)
+>>> dummy_input = generate_dummy_input(feature_info['features'], batch_size=2)
+"""
+
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+
+# Import feature types for type checking
+try:
+    from ..basic.features import DenseFeature, SequenceFeature, SparseFeature
+except ImportError:
+    # Fallback for standalone usage
+    SparseFeature = None
+    DenseFeature = None
+    SequenceFeature = None
+
+
+def extract_feature_info(model: nn.Module) -> Dict[str, Any]:
+    """Extract feature information from a torch-rechub model using reflection.
+
+    This function inspects model attributes to find feature lists without
+    modifying the model code. Supports various model architectures.
+
+    Parameters
+    ----------
+    model : nn.Module
+        The recommendation model to inspect.
+
+    Returns
+    -------
+    dict
+        Dictionary containing:
+        - 'features': List of unique Feature objects
+        - 'input_names': List of feature names in order
+        - 'input_types': Dict mapping feature name to feature type
+        - 'user_features': List of user-side features (for dual-tower models)
+        - 'item_features': List of item-side features (for dual-tower models)
+
+    Examples
+    --------
+    >>> from torch_rechub.models.ranking import DeepFM
+    >>> model = DeepFM(deep_features, fm_features, mlp_params)
+    >>> info = extract_feature_info(model)
+    >>> print(info['input_names'])  # ['user_id', 'item_id', ...]
+    """
+    # Common feature attribute names across different model types
+    feature_attrs = [
+        'features',           # MMOE, DCN, etc.
+        'deep_features',      # DeepFM, WideDeep
+        'fm_features',        # DeepFM
+        'wide_features',      # WideDeep
+        'linear_features',    # DeepFFM
+        'cross_features',     # DeepFFM
+        'user_features',      # DSSM, YoutubeDNN, MIND
+        'item_features',      # DSSM, YoutubeDNN, MIND
+        'history_features',   # DIN, MIND
+        'target_features',    # DIN
+        'neg_item_feature',   # YoutubeDNN, MIND
+    ]
+
+    all_features = []
+    user_features = []
+    item_features = []
+
+    for attr in feature_attrs:
+        if hasattr(model, attr):
+            feat_list = getattr(model, attr)
+            if isinstance(feat_list, list) and len(feat_list) > 0:
+                all_features.extend(feat_list)
+                # Track user/item features for dual-tower models
+                if 'user' in attr or 'history' in attr:
+                    user_features.extend(feat_list)
+                elif 'item' in attr:
+                    item_features.extend(feat_list)
+
+    # Deduplicate features by name while preserving order
+    seen = set()
+    unique_features = []
+    for f in all_features:
+        if hasattr(f, 'name') and f.name not in seen:
+            seen.add(f.name)
+            unique_features.append(f)
+
+    # Deduplicate user/item features
+    seen_user = set()
+    unique_user = [f for f in user_features if hasattr(f, 'name') and f.name not in seen_user and not seen_user.add(f.name)]
+    seen_item = set()
+    unique_item = [f for f in item_features if hasattr(f, 'name') and f.name not in seen_item and not seen_item.add(f.name)]
+
+    # Build input names and types
+    input_names = [f.name for f in unique_features if hasattr(f, 'name')]
+    input_types = {f.name: type(f).__name__ for f in unique_features if hasattr(f, 'name')}
+
+    return {
+        'features': unique_features,
+        'input_names': input_names,
+        'input_types': input_types,
+        'user_features': unique_user,
+        'item_features': unique_item,
+    }
+
+
+def generate_dummy_input(features: List[Any], batch_size: int = 2, seq_length: int = 10, device: str = 'cpu') -> Tuple[torch.Tensor, ...]:
+    """Generate dummy input tensors based on feature definitions.
+
+    Parameters
+    ----------
+    features : list
+        List of Feature objects (SparseFeature, DenseFeature, SequenceFeature).
+    batch_size : int, default=2
+        Batch size for dummy input.
+    seq_length : int, default=10
+        Sequence length for SequenceFeature.
+    device : str, default='cpu'
+        Device to create tensors on.
+
+    Returns
+    -------
+    tuple of Tensor
+        Tuple of tensors in the order of input features.
+
+    Examples
+    --------
+    >>> features = [SparseFeature("user_id", 1000), SequenceFeature("hist", 500)]
+    >>> dummy = generate_dummy_input(features, batch_size=4)
+    >>> # Returns (user_id_tensor[4], hist_tensor[4, 10])
+    """
+    # Dynamic import to handle feature types
+    from ..basic.features import DenseFeature, SequenceFeature, SparseFeature
+
+    inputs = []
+    for feat in features:
+        if isinstance(feat, SequenceFeature):
+            # Sequence features have shape [batch_size, seq_length]
+            tensor = torch.randint(0, feat.vocab_size, (batch_size, seq_length), device=device)
+        elif isinstance(feat, SparseFeature):
+            # Sparse features have shape [batch_size]
+            tensor = torch.randint(0, feat.vocab_size, (batch_size,), device=device)
+        elif isinstance(feat, DenseFeature):
+            # Dense features always have shape [batch_size, embed_dim]
+            tensor = torch.randn(batch_size, feat.embed_dim, device=device)
+        else:
+            raise TypeError(f"Unsupported feature type: {type(feat)}")
+        inputs.append(tensor)
+    return tuple(inputs)
+
+
+def generate_dummy_input_dict(features: List[Any], batch_size: int = 2, seq_length: int = 10, device: str = 'cpu') -> Dict[str, torch.Tensor]:
+    """Generate dummy input dict based on feature definitions.
+
+    Similar to generate_dummy_input but returns a dict mapping feature names
+    to tensors. This is the expected input format for torch-rechub models.
+
+    Parameters
+    ----------
+    features : list
+        List of Feature objects (SparseFeature, DenseFeature, SequenceFeature).
+    batch_size : int, default=2
+        Batch size for dummy input.
+    seq_length : int, default=10
+        Sequence length for SequenceFeature.
+    device : str, default='cpu'
+        Device to create tensors on.
+
+    Returns
+    -------
+    dict
+        Dict mapping feature names to tensors.
+
+    Examples
+    --------
+    >>> features = [SparseFeature("user_id", 1000)]
+    >>> dummy = generate_dummy_input_dict(features, batch_size=4)
+    >>> # Returns {"user_id": tensor[4]}
+    """
+    dummy_tuple = generate_dummy_input(features, batch_size, seq_length, device)
+    input_names = [f.name for f in features if hasattr(f, 'name')]
+    return {name: tensor for name, tensor in zip(input_names, dummy_tuple)}
+
+
+def generate_dynamic_axes(input_names: List[str], output_names: Optional[List[str]] = None, batch_dim: int = 0, include_seq_dim: bool = True, seq_features: Optional[List[str]] = None) -> Dict[str, Dict[int, str]]:
+    """Generate dynamic axes configuration for ONNX export.
+
+    Parameters
+    ----------
+    input_names : list of str
+        List of input tensor names.
+    output_names : list of str, optional
+        List of output tensor names. Default is ["output"].
+    batch_dim : int, default=0
+        Dimension index for batch size.
+    include_seq_dim : bool, default=True
+        Whether to include sequence dimension as dynamic.
+    seq_features : list of str, optional
+        List of feature names that are sequences.
+
+    Returns
+    -------
+    dict
+        Dynamic axes dict for torch.onnx.export.
+
+    Examples
+    --------
+    >>> axes = generate_dynamic_axes(["user_id", "item_id"], seq_features=["hist"])
+    >>> # Returns {"user_id": {0: "batch_size"}, "item_id": {0: "batch_size"}, ...}
+    """
+    if output_names is None:
+        output_names = ["output"]
+
+    dynamic_axes = {}
+
+    # Input axes
+    for name in input_names:
+        dynamic_axes[name] = {batch_dim: "batch_size"}
+        # Add sequence dimension for sequence features
+        if include_seq_dim and seq_features and name in seq_features:
+            dynamic_axes[name][1] = "seq_length"
+
+    # Output axes
+    for name in output_names:
+        dynamic_axes[name] = {batch_dim: "batch_size"}
+
+    return dynamic_axes

torch_rechub/utils/onnx_export.py

@@ -62,142 +62,9 @@ class ONNXWrapper(nn.Module):
             self.model.mode = self._original_mode
 
 
-def extract_feature_info(model: nn.Module) -> Dict[str, Any]:
-    """Extract feature information from a model using reflection.
-
-    This function inspects model attributes to find feature lists without
-    modifying the model code. Supports various model architectures.
-
-    Args:
-        model: The recommendation model to inspect.
-
-    Returns:
-        Dict containing:
-            - 'features': List of unique Feature objects
-            - 'input_names': List of feature names in order
-            - 'input_types': Dict mapping feature name to feature type
-            - 'user_features': List of user-side features (for dual-tower models)
-            - 'item_features': List of item-side features (for dual-tower models)
-    """
-    # Common feature attribute names across different model types
-    feature_attrs = [
-        'features',           # MMOE, DCN, etc.
-        'deep_features',      # DeepFM, WideDeep
-        'fm_features',        # DeepFM
-        'wide_features',      # WideDeep
-        'linear_features',    # DeepFFM
-        'cross_features',     # DeepFFM
-        'user_features',      # DSSM, YoutubeDNN, MIND
-        'item_features',      # DSSM, YoutubeDNN, MIND
-        'history_features',   # DIN, MIND
-        'target_features',    # DIN
-        'neg_item_feature',   # YoutubeDNN, MIND
-    ]
-
-    all_features = []
-    user_features = []
-    item_features = []
-
-    for attr in feature_attrs:
-        if hasattr(model, attr):
-            feat_list = getattr(model, attr)
-            if isinstance(feat_list, list) and len(feat_list) > 0:
-                all_features.extend(feat_list)
-                # Track user/item features for dual-tower models
-                if 'user' in attr or 'history' in attr:
-                    user_features.extend(feat_list)
-                elif 'item' in attr:
-                    item_features.extend(feat_list)
-
-    # Deduplicate features by name while preserving order
-    seen = set()
-    unique_features = []
-    for f in all_features:
-        if hasattr(f, 'name') and f.name not in seen:
-            seen.add(f.name)
-            unique_features.append(f)
-
-    # Deduplicate user/item features
-    seen_user = set()
-    unique_user = [f for f in user_features if hasattr(f, 'name') and f.name not in seen_user and not seen_user.add(f.name)]
-    seen_item = set()
-    unique_item = [f for f in item_features if hasattr(f, 'name') and f.name not in seen_item and not seen_item.add(f.name)]
-
-    return {
-        'features': unique_features,
-        'input_names': [f.name for f in unique_features],
-        'input_types': {
-            f.name: type(f).__name__ for f in unique_features
-        },
-        'user_features': unique_user,
-        'item_features': unique_item,
-    }
-
-
-def generate_dummy_input(features: List[Any], batch_size: int = 2, seq_length: int = 10, device: str = 'cpu') -> Tuple[torch.Tensor, ...]:
-    """Generate dummy input tensors for ONNX export based on feature definitions.
-
-    Args:
-        features: List of Feature objects (SparseFeature, DenseFeature, SequenceFeature).
-        batch_size: Batch size for dummy input (default: 2).
-        seq_length: Sequence length for SequenceFeature (default: 10).
-        device: Device to create tensors on (default: 'cpu').
-
-    Returns:
-        Tuple of tensors in the order of input features.
-
-    Example:
-        >>> features = [SparseFeature("user_id", 1000), SequenceFeature("hist", 500)]
-        >>> dummy = generate_dummy_input(features, batch_size=4)
-        >>> # Returns (user_id_tensor[4], hist_tensor[4, 10])
-    """
-    inputs = []
-    for feat in features:
-        if isinstance(feat, SequenceFeature):
-            # Sequence features have shape [batch_size, seq_length]
-            tensor = torch.randint(0, feat.vocab_size, (batch_size, seq_length), device=device)
-        elif isinstance(feat, SparseFeature):
-            # Sparse features have shape [batch_size]
-            tensor = torch.randint(0, feat.vocab_size, (batch_size,), device=device)
-        elif isinstance(feat, DenseFeature):
-            # Dense features have shape [batch_size, embed_dim]
-            tensor = torch.randn(batch_size, feat.embed_dim, device=device)
-        else:
-            raise TypeError(f"Unsupported feature type: {type(feat)}")
-        inputs.append(tensor)
-    return tuple(inputs)
-
-
-def generate_dynamic_axes(input_names: List[str], output_names: List[str] = None, batch_dim: int = 0, include_seq_dim: bool = True, seq_features: List[str] = None) -> Dict[str, Dict[int, str]]:
-    """Generate dynamic axes configuration for ONNX export.
-
-    Args:
-        input_names: List of input tensor names.
-        output_names: List of output tensor names (default: ["output"]).
-        batch_dim: Dimension index for batch size (default: 0).
-        include_seq_dim: Whether to include sequence dimension as dynamic (default: True).
-        seq_features: List of feature names that are sequences (default: auto-detect).
-
-    Returns:
-        Dynamic axes dict for torch.onnx.export.
-    """
-    if output_names is None:
-        output_names = ["output"]
-
-    dynamic_axes = {}
-
-    # Input axes
-    for name in input_names:
-        dynamic_axes[name] = {batch_dim: "batch_size"}
-        # Add sequence dimension for sequence features
-        if include_seq_dim and seq_features and name in seq_features:
-            dynamic_axes[name][1] = "seq_length"
-
-    # Output axes
-    for name in output_names:
-        dynamic_axes[name] = {batch_dim: "batch_size"}
-
-    return dynamic_axes
+# Re-export from model_utils for backward compatibility
+# The actual implementations are now in model_utils.py
+from .model_utils import extract_feature_info, generate_dummy_input, generate_dummy_input_dict, generate_dynamic_axes
 
 
 class ONNXExporter:

torch_rechub/utils/visualization.py

@@ -0,0 +1,271 @@
+"""
+Model Visualization Utilities for Torch-RecHub.
+
+This module provides model structure visualization using torchview library.
+Requires optional dependencies: pip install torch-rechub[visualization]
+
+Example:
+    >>> from torch_rechub.utils.visualization import visualize_model, display_graph
+    >>> graph = visualize_model(model, depth=4)
+    >>> display_graph(graph)  # Display in Jupyter Notebook
+
+    >>> # Save to file
+    >>> visualize_model(model, save_path="model_arch.pdf")
+"""
+
+from typing import Any, Dict, List, Optional, Union
+
+import torch
+import torch.nn as nn
+
+# Check for optional dependencies
+TORCHVIEW_AVAILABLE = False
+TORCHVIEW_SKIP_REASON = "torchview not installed"
+
+try:
+    from torchview import draw_graph
+    TORCHVIEW_AVAILABLE = True
+except ImportError as e:
+    TORCHVIEW_SKIP_REASON = f"torchview not available: {e}"
+
+
+def _is_jupyter_environment() -> bool:
+    """Check if running in Jupyter/IPython environment."""
+    try:
+        from IPython import get_ipython
+        shell = get_ipython()
+        if shell is None:
+            return False
+        # Check for Jupyter notebook or qtconsole
+        shell_class = shell.__class__.__name__
+        return shell_class in ('ZMQInteractiveShell', 'TerminalInteractiveShell')
+    except (ImportError, NameError):
+        return False
+
+
+def display_graph(graph: Any, format: str = 'png') -> Any:
+    """Display a torchview ComputationGraph in Jupyter Notebook.
+
+    This function provides a reliable way to display visualization graphs
+    in Jupyter environments, especially VSCode Jupyter.
+
+    Parameters
+    ----------
+    graph : ComputationGraph
+        A torchview ComputationGraph object returned by visualize_model().
+    format : str, default='png'
+        Output format, 'png' recommended for VSCode.
+
+    Returns
+    -------
+    graphviz.Digraph or None
+        The displayed graph object, or None if display fails.
+
+    Examples
+    --------
+    >>> graph = visualize_model(model, depth=4)
+    >>> display_graph(graph)  # Works in VSCode Jupyter
+    """
+    if not TORCHVIEW_AVAILABLE:
+        raise ImportError(f"Visualization requires torchview. {TORCHVIEW_SKIP_REASON}\n"
+                          "Install with: pip install torch-rechub[visualization]")
+
+    try:
+        import graphviz
+
+        # Set format for VSCode compatibility
+        graphviz.set_jupyter_format(format)
+    except ImportError:
+        pass
+
+    # Get the visual_graph (graphviz.Digraph object)
+    visual = graph.visual_graph
+
+    # Try to use IPython display for explicit rendering
+    try:
+        from IPython.display import display
+        display(visual)
+        return visual
+    except ImportError:
+        # Not in IPython/Jupyter environment, return the graph
+        return visual
+
+
+def visualize_model(
+    model: nn.Module,
+    input_data: Optional[Dict[str,
+                              torch.Tensor]] = None,
+    batch_size: int = 2,
+    seq_length: int = 10,
+    depth: int = 3,
+    show_shapes: bool = True,
+    expand_nested: bool = True,
+    save_path: Optional[str] = None,
+    graph_name: str = "model",
+    device: str = "cpu",
+    dpi: int = 300,
+    **kwargs
+) -> Any:
+    """Visualize a Torch-RecHub model's computation graph.
+
+    This function generates a visual representation of the model architecture,
+    showing layer connections, tensor shapes, and nested module structures.
+    It automatically extracts feature information from the model to generate
+    appropriate dummy inputs.
+
+    Parameters
+    ----------
+    model : nn.Module
+        PyTorch model to visualize. Should be a Torch-RecHub model
+        with feature attributes (e.g., DeepFM, DSSM, MMOE).
+    input_data : dict, optional
+        Dict of example inputs {feature_name: tensor}.
+        If None, inputs are auto-generated based on model features.
+    batch_size : int, default=2
+        Batch size for auto-generated inputs.
+    seq_length : int, default=10
+        Sequence length for SequenceFeature inputs.
+    depth : int, default=3
+        Visualization depth - higher values show more detail.
+        Set to -1 to show all layers.
+    show_shapes : bool, default=True
+        Whether to display tensor shapes on edges.
+    expand_nested : bool, default=True
+        Whether to expand nested nn.Module with dashed borders.
+    save_path : str, optional
+        Path to save the graph image. Supports .pdf, .svg, .png formats.
+        If None, displays in Jupyter or opens system viewer.
+    graph_name : str, default="model"
+        Name for the computation graph.
+    device : str, default="cpu"
+        Device for model execution during tracing.
+    dpi : int, default=300
+        Resolution in dots per inch for output image.
+        Higher values produce sharper images suitable for papers.
+    **kwargs : dict
+        Additional arguments passed to torchview.draw_graph().
+
+    Returns
+    -------
+    ComputationGraph
+        A torchview ComputationGraph object.
+        - Use `.visual_graph` property to get the graphviz.Digraph
+        - Use `.resize_graph(scale=1.5)` to adjust graph size
+
+    Raises
+    ------
+    ImportError
+        If torchview or graphviz is not installed.
+    ValueError
+        If model has no recognizable feature attributes.
+
+    Notes
+    -----
+    Default Display Behavior:
+        When `save_path` is None (default):
+        - In Jupyter/IPython: automatically displays the graph inline
+        - In Python script: opens the graph with system default viewer
+
+    Requires graphviz system package: apt/brew/choco install graphviz.
+    For Jupyter display issues, try: graphviz.set_jupyter_format('png').
+
+    Examples
+    --------
+    >>> from torch_rechub.models.ranking import DeepFM
+    >>> from torch_rechub.utils.visualization import visualize_model
+    >>>
+    >>> # Auto-display in Jupyter or open in viewer
+    >>> visualize_model(model, depth=4)  # No save_path needed
+    >>>
+    >>> # Save to high-DPI PNG for paper
+    >>> visualize_model(model, save_path="model.png", dpi=300)
+    """
+    if not TORCHVIEW_AVAILABLE:
+        raise ImportError(
+            f"Visualization requires torchview. {TORCHVIEW_SKIP_REASON}\n"
+            "Install with: pip install torch-rechub[visualization]\n"
+            "Also ensure graphviz is installed on your system:\n"
+            "  - Ubuntu/Debian: sudo apt-get install graphviz\n"
+            "  - macOS: brew install graphviz\n"
+            "  - Windows: choco install graphviz"
+        )
+
+    # Import feature extraction utilities from model_utils
+    from .model_utils import extract_feature_info, generate_dummy_input_dict
+
+    model.eval()
+    model.to(device)
+
+    # Auto-generate input data if not provided
+    if input_data is None:
+        feature_info = extract_feature_info(model)
+        features = feature_info['features']
+
+        if not features:
+            raise ValueError("Could not extract feature information from model. "
+                             "Please provide input_data parameter manually.")
+
+        # Generate dummy input dict
+        input_data = generate_dummy_input_dict(features, batch_size=batch_size, seq_length=seq_length, device=device)
+    else:
+        # Ensure input tensors are on correct device
+        input_data = {k: v.to(device) for k, v in input_data.items()}
+
+    # IMPORTANT: Wrap input_data dict in a tuple to work around torchview's behavior
+    #
+    # torchview's forward_prop function checks the type of input_data:
+    #   - If isinstance(x, (list, tuple)): model(*x)
+    #   - If isinstance(x, Mapping): model(**x)  <- This unpacks dict as kwargs!
+    #   - Else: model(x)
+    #
+    # torch-rechub models expect forward(self, x) where x is a complete dict.
+    # By wrapping the dict in a tuple, torchview will call:
+    #   model(*(input_dict,)) = model(input_dict)
+    # which is exactly what our models expect.
+    input_data_wrapped = (input_data,)
+
+    # Call torchview.draw_graph without saving (we'll save manually with DPI)
+    graph = draw_graph(
+        model,
+        input_data=input_data_wrapped,
+        graph_name=graph_name,
+        depth=depth,
+        device=device,
+        expand_nested=expand_nested,
+        show_shapes=show_shapes,
+        save_graph=False,  # Don't save here, we'll save manually with DPI
+        **kwargs
+    )
+
+    # Set DPI for high-quality output (must be set BEFORE rendering/saving)
+    graph.visual_graph.graph_attr['dpi'] = str(dpi)
+
+    # Handle save_path: manually save with DPI applied
+    if save_path:
+        import os
+        directory = os.path.dirname(save_path) or "."
+        filename = os.path.splitext(os.path.basename(save_path))[0]
+        ext = os.path.splitext(save_path)[1].lstrip('.')
+        # Default to pdf if no extension
+        output_format = ext if ext else 'pdf'
+        # Create directory if it doesn't exist
+        if directory != "." and not os.path.exists(directory):
+            os.makedirs(directory, exist_ok=True)
+        # Render and save with DPI applied
+        graph.visual_graph.render(
+            filename=filename,
+            directory=directory,
+            format=output_format,
+            cleanup=True  # Remove intermediate .gv file
+        )
+
+    # Handle default display behavior when save_path is None
+    if save_path is None:
+        if _is_jupyter_environment():
+            # In Jupyter: display inline
+            display_graph(graph)
+        else:
+            # In script: open with system viewer
+            graph.visual_graph.view(cleanup=True)
+
+    return graph