dragon-ml-toolbox 19.13.0__py3-none-any.whl → 20.0.0__py3-none-any.whl

ml_tools/_core/_ML_models.py

@@ -1,668 +0,0 @@
-import torch
-from torch import nn
-from typing import List, Union, Tuple, Dict, Any
-from pathlib import Path
-import json
-
-from ._logger import get_logger
-from ._path_manager import make_fullpath
-from ._script_info import _script_info
-from ._keys import PytorchModelArchitectureKeys
-from ._schema import FeatureSchema
-
-
-_LOGGER = get_logger("DragonModel")
-
-
-__all__ = [
-    "DragonMLP",
-    "DragonAttentionMLP",
-    "DragonMultiHeadAttentionNet",
-    "DragonTabularTransformer"
-]
-
-
-class _ArchitectureHandlerMixin:
-    """
-    A mixin class to provide save and load functionality for model architectures.
-    """
-    def save(self: nn.Module, directory: Union[str, Path], verbose: bool = True): # type: ignore
-        """Saves the model's architecture to a JSON file."""
-        if not hasattr(self, 'get_architecture_config'):
-            _LOGGER.error(f"Model '{self.__class__.__name__}' must have a 'get_architecture_config()' method to use this functionality.")
-            raise AttributeError()
-
-        path_dir = make_fullpath(directory, make=True, enforce="directory")
-        
-        json_filename = PytorchModelArchitectureKeys.SAVENAME + ".json"
-        
-        full_path = path_dir / json_filename
-
-        config = {
-            PytorchModelArchitectureKeys.MODEL: self.__class__.__name__,
-            PytorchModelArchitectureKeys.CONFIG: self.get_architecture_config() # type: ignore
-        }
-
-        with open(full_path, 'w') as f:
-            json.dump(config, f, indent=4)
-        
-        if verbose:
-            _LOGGER.info(f"Architecture for '{self.__class__.__name__}' saved as '{full_path.name}'")
-
-    @classmethod
-    def load(cls: type, file_or_dir: Union[str, Path], verbose: bool = True) -> nn.Module:
-        """Loads a model architecture from a JSON file. If a directory is provided, the function will attempt to load a JSON file inside."""
-        user_path = make_fullpath(file_or_dir)
-        
-        if user_path.is_dir():
-            json_filename = PytorchModelArchitectureKeys.SAVENAME + ".json"
-            target_path = make_fullpath(user_path / json_filename, enforce="file")
-        elif user_path.is_file():
-            target_path = user_path
-        else:
-            _LOGGER.error(f"Invalid path: '{file_or_dir}'")
-            raise IOError()
-
-        with open(target_path, 'r') as f:
-            saved_data = json.load(f)
-
-        saved_class_name = saved_data[PytorchModelArchitectureKeys.MODEL]
-        config = saved_data[PytorchModelArchitectureKeys.CONFIG]
-
-        if saved_class_name != cls.__name__:
-            _LOGGER.error(f"Model class mismatch. File specifies '{saved_class_name}', but '{cls.__name__}' was expected.")
-            raise ValueError()
-
-        model = cls(**config)
-        if verbose:
-            _LOGGER.info(f"Successfully loaded architecture for '{saved_class_name}'")
-        return model
-
-
-class _BaseMLP(nn.Module, _ArchitectureHandlerMixin):
-    """
-    A base class for Multilayer Perceptrons.
-    
-    Handles validation, configuration, and the creation of the core MLP layers,
-    allowing subclasses to define their own pre-processing and forward pass.
-    """
-    def __init__(self, 
-                 in_features: int, 
-                 out_targets: int,
-                 hidden_layers: List[int], 
-                 drop_out: float) -> None:
-        super().__init__()
-
-        # --- Validation ---
-        if not isinstance(in_features, int) or in_features < 1:
-            _LOGGER.error("'in_features' must be a positive integer.")
-            raise ValueError()
-        if not isinstance(out_targets, int) or out_targets < 1:
-            _LOGGER.error("'out_targets' must be a positive integer.")
-            raise ValueError()
-        if not isinstance(hidden_layers, list) or not all(isinstance(n, int) for n in hidden_layers):
-            _LOGGER.error("'hidden_layers' must be a list of integers.")
-            raise TypeError()
-        if not (0.0 <= drop_out < 1.0):
-            _LOGGER.error("'drop_out' must be a float between 0.0 and 1.0.")
-            raise ValueError()
-        
-        # --- Save configuration ---
-        self.in_features = in_features
-        self.out_targets = out_targets
-        self.hidden_layers = hidden_layers
-        self.drop_out = drop_out
-
-        # --- Build the core MLP network ---
-        mlp_layers = []
-        current_features = in_features
-        for neurons in hidden_layers:
-            mlp_layers.extend([
-                nn.Linear(current_features, neurons),
-                nn.BatchNorm1d(neurons),
-                nn.ReLU(),
-                nn.Dropout(p=drop_out)
-            ])
-            current_features = neurons
-        
-        self.mlp = nn.Sequential(*mlp_layers)
-        # Set a customizable Prediction Head for flexibility, specially in transfer learning and fine-tuning
-        self.output_layer = nn.Linear(current_features, out_targets)
-
-    def get_architecture_config(self) -> Dict[str, Any]:
-        """Returns the base configuration of the model."""
-        return {
-            'in_features': self.in_features,
-            'out_targets': self.out_targets,
-            'hidden_layers': self.hidden_layers,
-            'drop_out': self.drop_out
-        }
-        
-    def _repr_helper(self, name: str, mlp_layers: list[str]):
-        last_layer = self.output_layer
-        if isinstance(last_layer, nn.Linear):
-            mlp_layers.append(str(last_layer.out_features))
-        else:
-            mlp_layers.append("Custom Prediction Head")
-        
-        # Creates a string like: 10 -> 40 -> 80 -> 40 -> 2
-        arch_str = ' -> '.join(mlp_layers)
-        
-        return f"{name}(arch: {arch_str})"
-
-
-class _BaseAttention(_BaseMLP):
-    """
-    Abstract base class for MLP models that incorporate an attention mechanism
-    before the main MLP layers.
-    """
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        # By default, models inheriting this do not have the flag.
-        self.attention = None
-        self.has_interpretable_attention = False
-    
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """Defines the standard forward pass."""
-        logits, _attention_weights = self.forward_attention(x)
-        return logits
-
-    def forward_attention(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
-        """Returns logits and attention weights."""
-        # This logic is now shared and defined in one place
-        x, attention_weights = self.attention(x) # type: ignore
-        x = self.mlp(x)
-        logits = self.output_layer(x)
-        return logits, attention_weights
-
-
-class DragonMLP(_BaseMLP):
-    """
-    Creates a versatile Multilayer Perceptron (MLP) for regression or classification tasks.
-    """
-    def __init__(self, in_features: int, out_targets: int,
-                 hidden_layers: List[int] = [256, 128], drop_out: float = 0.2) -> None:
-        """
-        Args:
-            in_features (int): The number of input features (e.g., columns in your data).
-            out_targets (int): The number of output targets. For regression, this is
-                typically 1. For classification, it's the number of classes.
-            hidden_layers (list[int]): A list where each integer represents the
-                number of neurons in a hidden layer.
-            drop_out (float): The dropout probability for neurons in each hidden
-                layer. Must be between 0.0 and 1.0.
-                
-        ### Rules of thumb:
-        - Choose a number of hidden neurons between the size of the input layer and the size of the output layer. 
-        - The number of hidden neurons should be 2/3 the size of the input layer, plus the size of the output layer. 
-        - The number of hidden neurons should be less than twice the size of the input layer.
-        """
-        super().__init__(in_features, out_targets, hidden_layers, drop_out)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """Defines the forward pass of the model."""
-        x = self.mlp(x)
-        logits = self.output_layer(x)
-        return logits
-    
-    def __repr__(self) -> str:
-        """Returns the developer-friendly string representation of the model."""
-        # Extracts the number of neurons from each nn.Linear layer
-        layer_sizes = [str(layer.in_features) for layer in self.mlp if isinstance(layer, nn.Linear)]
-        
-        return self._repr_helper(name="DragonMLP", mlp_layers=layer_sizes)
-
-
-class DragonAttentionMLP(_BaseAttention):
-    """
-    A Multilayer Perceptron (MLP) that incorporates an Attention layer to dynamically weigh input features.
-    
-    In inference mode use `forward_attention()` to get a tuple with `(output, attention_weights)`
-    """
-    def __init__(self, in_features: int, out_targets: int,
-                 hidden_layers: List[int] = [256, 128], drop_out: float = 0.2) -> None:
-        """
-        Args:
-            in_features (int): The number of input features (e.g., columns in your data).
-            out_targets (int): The number of output targets. For regression, this is
-                typically 1. For classification, it's the number of classes.
-            hidden_layers (list[int]): A list where each integer represents the
-                number of neurons in a hidden layer.
-            drop_out (float): The dropout probability for neurons in each hidden
-                layer. Must be between 0.0 and 1.0.
-        """
-        super().__init__(in_features, out_targets, hidden_layers, drop_out)
-        # Attention
-        self.attention = _AttentionLayer(in_features)
-        self.has_interpretable_attention = True
-    
-    def __repr__(self) -> str:
-        """Returns the developer-friendly string representation of the model."""
-        # Start with the input features and the attention marker
-        arch = [str(self.in_features), "[Attention]"]
-
-        # Find all other linear layers in the MLP 
-        for layer in self.mlp[1:]:
-            if isinstance(layer, nn.Linear):
-                arch.append(str(layer.in_features))
-        
-        return self._repr_helper(name="DragonAttentionMLP", mlp_layers=arch)
-
-
-class DragonMultiHeadAttentionNet(_BaseAttention):
-    """
-    An MLP that incorporates a standard `nn.MultiheadAttention` layer to process
-    the input features.
-
-    In inference mode use `forward_attention()` to get a tuple with `(output, attention_weights)`.
-    """
-    def __init__(self, in_features: int, out_targets: int,
-                 hidden_layers: List[int] = [256, 128], drop_out: float = 0.2,
-                 num_heads: int = 4, attention_dropout: float = 0.1) -> None:
-        """
-        Args:
-            in_features (int): The number of input features.
-            out_targets (int): The number of output targets.
-            hidden_layers (list[int]): A list of neuron counts for each hidden layer.
-            drop_out (float): The dropout probability for the MLP layers.
-            num_heads (int): The number of attention heads.
-            attention_dropout (float): Dropout probability in the attention layer.
-        """
-        super().__init__(in_features, out_targets, hidden_layers, drop_out)
-        self.num_heads = num_heads
-        self.attention_dropout = attention_dropout
-        
-        self.attention = _MultiHeadAttentionLayer(
-            num_features=in_features,
-            num_heads=num_heads,
-            dropout=attention_dropout
-        )
-
-    def get_architecture_config(self) -> Dict[str, Any]:
-        """Returns the full configuration of the model."""
-        config = super().get_architecture_config()
-        config['num_heads'] = self.num_heads
-        config['attention_dropout'] = self.attention_dropout
-        return config
-    
-    def __repr__(self) -> str:
-        """Returns the developer-friendly string representation of the model."""
-        mlp_part = " -> ".join(
-            [str(self.in_features)] + 
-            [str(h) for h in self.hidden_layers] + 
-            [str(self.out_targets)]
-        )
-        arch_str = f"{self.in_features} -> [MultiHead(h={self.num_heads})] -> {mlp_part}"
-        
-        return f"DragonMultiHeadAttentionNet(arch: {arch_str})"
-
-
-class DragonTabularTransformer(nn.Module, _ArchitectureHandlerMixin):
-    """
-    A Transformer-based model for tabular data tasks.
-    
-    This model uses a Feature Tokenizer to convert all input features into a
-    sequence of embeddings, prepends a [CLS] token, and processes the
-    sequence with a standard Transformer Encoder.
-    """
-    def __init__(self, *,
-                 schema: FeatureSchema,
-                 out_targets: int,
-                 embedding_dim: int = 256,
-                 num_heads: int = 8,
-                 num_layers: int = 6,
-                 dropout: float = 0.2):
-        """
-        Args:
-            schema (FeatureSchema): 
-                The definitive schema object created by `data_exploration.finalize_feature_schema()`.
-            out_targets (int): 
-                Number of output targets (1 for regression).
-            embedding_dim (int): 
-                The dimension for all feature embeddings. Must be divisible by num_heads. Common values: (64, 128, 192, 256, etc.)
-            num_heads (int): 
-                The number of heads in the multi-head attention mechanism. Common values: (4, 8, 16)
-            num_layers (int): 
-                The number of sub-encoder-layers in the transformer encoder. Common values: (4, 8, 12)
-            dropout (float): 
-                The dropout value.
-                
-        ## Note:
-        
-        **Embedding Dimension:** "Width" of the model. It's the N-dimension vector that will be used to represent each one of the features.
-            - Each continuous feature gets its own learnable N-dimension vector.
-            - Each categorical feature gets an embedding table that maps every category (e.g., "color=red", "color=blue") to a unique N-dimension vector.
-            
-        **Attention Heads:** Controls the "Multi-Head Attention" mechanism. Instead of looking at all the feature interactions at once, the model splits its attention into N parallel heads.
-            - Embedding Dimensions get divided by the number of Attention Heads, resulting in the dimensions assigned per head.
-
-        **Number of Layers:** "Depth" of the model. Number of identical `TransformerEncoderLayer` blocks that are stacked on top of each other.
-            - Layer 1: The attention heads find simple, direct interactions between the features.
-            - Layer 2: Takes the output of Layer 1 and finds interactions between those interactions and so on.
-            - Trade-off: More layers are more powerful but are slower to train and more prone to overfitting. If the training loss goes down but the validation loss goes up, you might have too many layers (or need more dropout).
-            
-        """
-        super().__init__()
-        
-         # --- Get info from schema ---
-        in_features = len(schema.feature_names)
-        categorical_index_map = schema.categorical_index_map
-
-         # --- Validation ---
-        if categorical_index_map and (max(categorical_index_map.keys()) >= in_features):
-            _LOGGER.error(f"A categorical index ({max(categorical_index_map.keys())}) is out of bounds for the provided input features ({in_features}).")
-            raise ValueError()
-        
-        # --- Save configuration ---
-        self.schema = schema # <-- Save the whole schema
-        self.out_targets = out_targets
-        self.embedding_dim = embedding_dim
-        self.num_heads = num_heads
-        self.num_layers = num_layers
-        self.dropout = dropout
-
-        # --- 1. Feature Tokenizer (now takes the schema) ---
-        self.tokenizer = _FeatureTokenizer(
-            schema=schema,
-            embedding_dim=embedding_dim
-        )
-        
-        # --- 2. CLS Token ---
-        self.cls_token = nn.Parameter(torch.randn(1, 1, embedding_dim))
-        
-        # --- 3. Transformer Encoder ---
-        encoder_layer = nn.TransformerEncoderLayer(
-            d_model=embedding_dim,
-            nhead=num_heads,
-            dropout=dropout,
-            batch_first=True # Crucial for (batch, seq, feature) input
-        )
-        self.transformer_encoder = nn.TransformerEncoder(
-            encoder_layer=encoder_layer,
-            num_layers=num_layers
-        )
-        
-        # --- 4. Prediction Head ---
-        self.output_layer = nn.Linear(embedding_dim, out_targets)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """Defines the forward pass of the model."""
-        # Get the batch size for later use
-        batch_size = x.shape[0]
-        
-        # 1. Get feature tokens from the tokenizer
-        # -> tokens shape: (batch_size, num_features, embedding_dim)
-        tokens = self.tokenizer(x)
-        
-        # 2. Prepend the [CLS] token to the sequence
-        # -> cls_tokens shape: (batch_size, 1, embedding_dim)
-        cls_tokens = self.cls_token.expand(batch_size, -1, -1)
-        # -> full_sequence shape: (batch_size, num_features + 1, embedding_dim)
-        full_sequence = torch.cat([cls_tokens, tokens], dim=1)
-
-        # 3. Pass the full sequence through the Transformer Encoder
-        # -> transformer_out shape: (batch_size, num_features + 1, embedding_dim)
-        transformer_out = self.transformer_encoder(full_sequence)
-        
-        # 4. Isolate the output of the [CLS] token (it's the first one)
-        # -> cls_output shape: (batch_size, embedding_dim)
-        cls_output = transformer_out[:, 0]
-        
-        # 5. Pass the [CLS] token's output through the prediction head
-        # -> logits shape: (batch_size, out_targets)
-        logits = self.output_layer(cls_output)
-        
-        return logits
-    
-    def get_architecture_config(self) -> Dict[str, Any]:
-        """Returns the full configuration of the model."""
-        # Deconstruct schema into a JSON-friendly dict
-        # Tuples are saved as lists
-        schema_dict = {
-            'feature_names': self.schema.feature_names,
-            'continuous_feature_names': self.schema.continuous_feature_names,
-            'categorical_feature_names': self.schema.categorical_feature_names,
-            'categorical_index_map': self.schema.categorical_index_map,
-            'categorical_mappings': self.schema.categorical_mappings
-        }
-
-        return {
-            'schema_dict': schema_dict,
-            'out_targets': self.out_targets,
-            'embedding_dim': self.embedding_dim,
-            'num_heads': self.num_heads,
-            'num_layers': self.num_layers,
-            'dropout': self.dropout
-        }
-    
-    @classmethod
-    def load(cls: type, file_or_dir: Union[str, Path], verbose: bool = True) -> nn.Module:
-        """Loads a model architecture from a JSON file."""
-        user_path = make_fullpath(file_or_dir)
-        
-        if user_path.is_dir():
-            json_filename = PytorchModelArchitectureKeys.SAVENAME + ".json"
-            target_path = make_fullpath(user_path / json_filename, enforce="file")
-        elif user_path.is_file():
-            target_path = user_path
-        else:
-            _LOGGER.error(f"Invalid path: '{file_or_dir}'")
-            raise IOError()
-
-        with open(target_path, 'r') as f:
-            saved_data = json.load(f)
-
-        saved_class_name = saved_data[PytorchModelArchitectureKeys.MODEL]
-        config = saved_data[PytorchModelArchitectureKeys.CONFIG]
-
-        if saved_class_name != cls.__name__:
-            _LOGGER.error(f"Model class mismatch. File specifies '{saved_class_name}', but '{cls.__name__}' was expected.")
-            raise ValueError()
-
-        # --- RECONSTRUCTION LOGIC ---
-        if 'schema_dict' not in config:
-            _LOGGER.error("Invalid architecture file: missing 'schema_dict'. This file may be from an older version.")
-            raise ValueError("Missing 'schema_dict' in config.")
-            
-        schema_data = config.pop('schema_dict')
-        
-        # Re-hydrate the categorical_index_map
-        # JSON saves all dict keys as strings, so we must convert them back to int.
-        raw_index_map = schema_data['categorical_index_map']
-        if raw_index_map is not None:
-            rehydrated_index_map = {int(k): v for k, v in raw_index_map.items()}
-        else:
-            rehydrated_index_map = None
-
-        # Re-hydrate the FeatureSchema object
-        # JSON deserializes tuples as lists, so we must convert them back.
-        schema = FeatureSchema(
-            feature_names=tuple(schema_data['feature_names']),
-            continuous_feature_names=tuple(schema_data['continuous_feature_names']),
-            categorical_feature_names=tuple(schema_data['categorical_feature_names']),
-            categorical_index_map=rehydrated_index_map,
-            categorical_mappings=schema_data['categorical_mappings']
-        )
-        
-        config['schema'] = schema
-        # --- End Reconstruction ---
-
-        model = cls(**config)
-        if verbose:
-            _LOGGER.info(f"Successfully loaded architecture for '{saved_class_name}'")
-        return model
-        
-    def __repr__(self) -> str:
-        """Returns the developer-friendly string representation of the model."""
-        # Build the architecture string part-by-part
-        parts = [
-            f"Tokenizer(features={len(self.schema.feature_names)}, dim={self.embedding_dim})",
-            "[CLS]",
-            f"TransformerEncoder(layers={self.num_layers}, heads={self.num_heads})",
-            f"PredictionHead(outputs={self.out_targets})"
-        ]
-        
-        arch_str = " -> ".join(parts)
-        
-        return f"DragonTabularTransformer(arch: {arch_str})"
-
-
-class _FeatureTokenizer(nn.Module):
-    """
-    Transforms raw numerical and categorical features from any column order 
-    into a sequence of embeddings.
-    """
-    def __init__(self,
-                 schema: FeatureSchema,
-                 embedding_dim: int):
-        """
-        Args:
-            schema (FeatureSchema): 
-                The definitive schema object from data_exploration.
-            embedding_dim (int): 
-                The dimension for all feature embeddings.
-        """
-        super().__init__()
-        
-        # --- Get info from schema ---
-        categorical_map = schema.categorical_index_map
-        
-        if categorical_map:
-            # Unpack the dictionary into separate lists
-            self.categorical_indices = list(categorical_map.keys())
-            cardinalities = list(categorical_map.values())
-        else:
-            self.categorical_indices = []
-            cardinalities = []
-        
-        # Derive numerical indices by finding what's not categorical
-        all_indices = set(range(len(schema.feature_names)))
-        categorical_indices_set = set(self.categorical_indices)
-        self.numerical_indices = sorted(list(all_indices - categorical_indices_set))
-        
-        self.embedding_dim = embedding_dim
-        
-        # A learnable embedding for each numerical feature
-        self.numerical_embeddings = nn.Parameter(torch.randn(len(self.numerical_indices), embedding_dim))
-        
-        # A standard embedding layer for each categorical feature
-        self.categorical_embeddings = nn.ModuleList(
-            [nn.Embedding(num_embeddings=c, embedding_dim=embedding_dim) for c in cardinalities]
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Processes features from a single input tensor and concatenates them
-        into a sequence of tokens.
-        """
-        # Select the correct columns for each type using the stored indices
-        x_numerical = x[:, self.numerical_indices].float()
-        x_categorical = x[:, self.categorical_indices].long()
-
-        # Process numerical features
-        numerical_tokens = x_numerical.unsqueeze(-1) * self.numerical_embeddings
-        
-        # Process categorical features
-        categorical_tokens = []
-        for i, embed_layer in enumerate(self.categorical_embeddings):
-            # x_categorical[:, i] selects the i-th categorical column
-            # (e.g., all values for the 'color' feature)
-            token = embed_layer(x_categorical[:, i]).unsqueeze(1)
-            categorical_tokens.append(token)
-        
-        # Concatenate all tokens into a single sequence
-        if not self.categorical_indices:
-             all_tokens = numerical_tokens
-        elif not self.numerical_indices:
-             all_tokens = torch.cat(categorical_tokens, dim=1)
-        else:
-             all_categorical_tokens = torch.cat(categorical_tokens, dim=1)
-             all_tokens = torch.cat([numerical_tokens, all_categorical_tokens], dim=1)
-        
-        return all_tokens
-
-
-class _AttentionLayer(nn.Module):
-    """
-    Calculates attention weights and applies them to the input features, incorporating a residual connection for improved stability and performance.
-    
-    Returns both the final output and the weights for interpretability.
-    """
-    def __init__(self, num_features: int):
-        super().__init__()
-        # The hidden layer size is a hyperparameter
-        hidden_size = max(16, num_features // 4)
-        
-        # Learn to produce attention scores
-        self.attention_net = nn.Sequential(
-            nn.Linear(num_features, hidden_size),
-            nn.Tanh(),
-            nn.Linear(hidden_size, num_features) # Output one score per feature
-        )
-        self.softmax = nn.Softmax(dim=1)
-
-    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        # x shape: (batch_size, num_features)
-        
-        # Get one raw "importance" score per feature
-        attention_scores = self.attention_net(x)
-        
-        # Apply the softmax module to get weights that sum to 1
-        attention_weights = self.softmax(attention_scores)
-        
-        # Weighted features (attention mechanism's output)
-        weighted_features = x * attention_weights
-        
-        # Residual connection
-        residual_connection = x + weighted_features
-        
-        return residual_connection, attention_weights
-
-
-class _MultiHeadAttentionLayer(nn.Module):
-    """
-    A wrapper for the standard `torch.nn.MultiheadAttention` layer.
-
-    This layer treats the entire input feature vector as a single item in a
-    sequence and applies self-attention to it. It is followed by a residual
-    connection and layer normalization, which is a standard block in
-    Transformer-style models.
-    """
-    def __init__(self, num_features: int, num_heads: int, dropout: float):
-        super().__init__()
-        self.attention = nn.MultiheadAttention(
-            embed_dim=num_features,
-            num_heads=num_heads,
-            dropout=dropout,
-            batch_first=True  # Crucial for (batch, seq, feature) input
-        )
-        self.layer_norm = nn.LayerNorm(num_features)
-
-    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        # x shape: (batch_size, num_features)
-
-        # nn.MultiheadAttention expects a sequence dimension.
-        # We add a sequence dimension of length 1.
-        # x_reshaped shape: (batch_size, 1, num_features)
-        x_reshaped = x.unsqueeze(1)
-
-        # Apply self-attention. query, key, and value are all the same.
-        # attn_output shape: (batch_size, 1, num_features)
-        # attn_weights shape: (batch_size, 1, 1)
-        attn_output, attn_weights = self.attention(
-            query=x_reshaped,
-            key=x_reshaped,
-            value=x_reshaped,
-            need_weights=True,
-            average_attn_weights=True # Average weights across heads
-        )
-
-        # Add residual connection and apply layer normalization (Post-LN)
-        out = self.layer_norm(x + attn_output.squeeze(1))
-
-        # Squeeze weights for a consistent output shape
-        return out, attn_weights.squeeze()
-
-
-def info():
-    _script_info(__all__)