Trixy3/trixy_core/ml/voice_recognition/model_architectures.py

"""
Model architectures for Voice Recognition system.

This module implements state-of-the-art speaker recognition model architectures:
- ECAPA-TDNN: Enhanced version of Time Delay Neural Networks with attention
- TitaNet-S: Small variant of the TitaNet architecture
- SpeakerNet-M: Medium-sized SpeakerNet architecture

All models are optimized for real-time inference and support dynamic speaker numbers.
"""

import logging
import math
from typing import Optional, List, Dict, Any, Tuple, Union
from enum import Enum
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import Parameter

# Configure logging
logger = logging.getLogger(__name__)


class ModelArchitectureType(Enum):
    """Supported model architecture types."""
    ECAPA_TDNN = "ecapa_tdnn"
    TITANET_S = "titanet_s"
    SPEAKERNET_M = "speakernet_m"


class TDNNBlock(nn.Module):
    """Time Delay Neural Network block."""
    
    def __init__(self, 
                 in_channels: int,
                 out_channels: int,
                 kernel_size: int,
                 dilation: int = 1,
                 dropout: float = 0.1,
                 use_batch_norm: bool = True):
        super().__init__()
        
        self.conv = nn.Conv1d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            dilation=dilation,
            padding=(kernel_size - 1) * dilation // 2
        )
        
        self.use_batch_norm = use_batch_norm
        if use_batch_norm:
            self.batch_norm = nn.BatchNorm1d(out_channels)
        
        self.activation = nn.ReLU()
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass through TDNN block."""
        x = self.conv(x)
        
        if self.use_batch_norm:
            x = self.batch_norm(x)
        
        x = self.activation(x)
        x = self.dropout(x)
        
        return x


class SEBlock(nn.Module):
    """Squeeze-and-Excitation block for channel attention."""
    
    def __init__(self, channels: int, reduction: int = 16):
        super().__init__()
        
        self.global_pool = nn.AdaptiveAvgPool1d(1)
        
        self.fc = nn.Sequential(
            nn.Linear(channels, channels // reduction),
            nn.ReLU(inplace=True),
            nn.Linear(channels // reduction, channels),
            nn.Sigmoid()
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass through SE block."""
        batch_size, channels, _ = x.size()
        
        # Global average pooling
        y = self.global_pool(x).view(batch_size, channels)
        
        # Squeeze and excitation
        y = self.fc(y).view(batch_size, channels, 1)
        
        # Scale input
        return x * y


class ResidualTDNNBlock(nn.Module):
    """Residual TDNN block with skip connections."""
    
    def __init__(self,
                 channels: int,
                 kernel_size: int,
                 dilation: int = 1,
                 dropout: float = 0.1,
                 use_se: bool = True):
        super().__init__()
        
        self.tdnn1 = TDNNBlock(channels, channels, kernel_size, dilation, dropout)
        self.tdnn2 = TDNNBlock(channels, channels, kernel_size, dilation, dropout)
        
        self.use_se = use_se
        if use_se:
            self.se = SEBlock(channels)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass through residual TDNN block."""
        residual = x
        
        x = self.tdnn1(x)
        x = self.tdnn2(x)
        
        if self.use_se:
            x = self.se(x)
        
        return x + residual


class AttentiveStatisticsPooling(nn.Module):
    """Attentive statistics pooling for temporal aggregation."""
    
    def __init__(self, input_dim: int):
        super().__init__()
        
        self.attention = nn.Linear(input_dim, 1)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass through attentive pooling.
        
        Args:
            x: Input tensor [batch_size, channels, time]
            
        Returns:
            Pooled tensor [batch_size, channels * 2]
        """
        # x: [batch_size, channels, time]
        batch_size, channels, time_steps = x.size()
        
        # Compute attention weights
        # Transpose to [batch_size, time, channels] for linear layer
        x_t = x.transpose(1, 2)  # [batch_size, time, channels]
        
        # Compute attention scores
        alpha = self.attention(x_t)  # [batch_size, time, 1]
        alpha = F.softmax(alpha, dim=1)  # [batch_size, time, 1]
        
        # Apply attention to get weighted mean
        mean = torch.sum(alpha * x_t, dim=1)  # [batch_size, channels]
        
        # Compute weighted standard deviation
        # First, compute weighted variance
        x_centered = x_t - mean.unsqueeze(1)  # [batch_size, time, channels]
        var = torch.sum(alpha * (x_centered ** 2), dim=1)  # [batch_size, channels]
        std = torch.sqrt(var + 1e-8)  # [batch_size, channels]
        
        # Concatenate mean and std
        pooled = torch.cat([mean, std], dim=1)  # [batch_size, channels * 2]
        
        return pooled


class ECADATDNNModel(nn.Module):
    """
    ECAPA-TDNN (Emphasized Channel Attention, Propagation and Aggregation TDNN).
    
    Enhanced version of TDNN with:
    - Squeeze-and-Excitation blocks for channel attention
    - Residual connections
    - Attentive statistics pooling
    """
    
    def __init__(self,
                 input_dim: int = 40,
                 embedding_dim: int = 192,
                 channels: int = 512,
                 num_speakers: int = 4,
                 dropout: float = 0.1,
                 use_attention_pooling: bool = True):
        super().__init__()
        
        self.input_dim = input_dim
        self.embedding_dim = embedding_dim
        self.num_speakers = num_speakers
        
        # Frame-level processing
        self.frame_layer = TDNNBlock(input_dim, channels, kernel_size=5, dropout=dropout)
        
        # TDNN blocks with increasing dilation
        self.tdnn_blocks = nn.ModuleList([
            ResidualTDNNBlock(channels, kernel_size=3, dilation=1, dropout=dropout),
            ResidualTDNNBlock(channels, kernel_size=3, dilation=2, dropout=dropout),
            ResidualTDNNBlock(channels, kernel_size=3, dilation=3, dropout=dropout),
            ResidualTDNNBlock(channels, kernel_size=3, dilation=1, dropout=dropout),
        ])
        
        # Segment-level processing
        self.segment_layer = TDNNBlock(channels, channels * 3, kernel_size=1, dropout=dropout)
        
        # Pooling layer
        self.use_attention_pooling = use_attention_pooling
        if use_attention_pooling:
            self.pooling = AttentiveStatisticsPooling(channels * 3)
            pooling_output_dim = channels * 3 * 2  # mean + std
        else:
            self.pooling = None
            pooling_output_dim = channels * 3 * 2  # manual mean + std
        
        # Embedding layer
        self.embedding_layer = nn.Sequential(
            nn.Linear(pooling_output_dim, embedding_dim),
            nn.BatchNorm1d(embedding_dim),
            nn.ReLU(),
            nn.Dropout(dropout)
        )
        
        # Classification layer (for training)
        self.classifier = nn.Linear(embedding_dim, num_speakers)
        
        # Initialize weights
        self._initialize_weights()
    
    def _initialize_weights(self):
        """Initialize model weights."""
        for m in self.modules():
            if isinstance(m, nn.Conv1d):
                nn.init.xavier_normal_(m.weight)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.xavier_normal_(m.weight)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm1d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
    
    def forward(self, x: torch.Tensor, extract_embedding: bool = False) -> torch.Tensor:
        """
        Forward pass through ECAPA-TDNN.
        
        Args:
            x: Input features [batch_size, input_dim, time]
            extract_embedding: Whether to return embeddings instead of logits
            
        Returns:
            Either speaker embeddings or classification logits
        """
        # Frame-level processing
        x = self.frame_layer(x)
        
        # TDNN blocks
        for tdnn_block in self.tdnn_blocks:
            x = tdnn_block(x)
        
        # Segment-level processing
        x = self.segment_layer(x)
        
        # Temporal pooling
        if self.use_attention_pooling:
            x = self.pooling(x)
        else:
            # Manual statistics pooling
            mean = torch.mean(x, dim=2)
            std = torch.std(x, dim=2)
            x = torch.cat([mean, std], dim=1)
        
        # Embedding extraction
        embeddings = self.embedding_layer(x)
        
        if extract_embedding:
            return F.normalize(embeddings, p=2, dim=1)
        
        # Classification
        logits = self.classifier(embeddings)
        return logits
    
    def extract_embedding(self, x: torch.Tensor) -> torch.Tensor:
        """Extract speaker embeddings."""
        return self.forward(x, extract_embedding=True)


class TitaNetSModel(nn.Module):
    """
    TitaNet-S (Small variant) architecture.
    
    Lightweight architecture optimized for real-time inference while
    maintaining good speaker recognition performance.
    """
    
    def __init__(self,
                 input_dim: int = 40,
                 embedding_dim: int = 192,
                 channels: List[int] = None,
                 num_speakers: int = 4,
                 dropout: float = 0.05):
        super().__init__()
        
        if channels is None:
            channels = [128, 256, 512, 1024]
        
        self.input_dim = input_dim
        self.embedding_dim = embedding_dim
        self.num_speakers = num_speakers
        
        # Input projection
        self.input_conv = nn.Conv1d(input_dim, channels[0], kernel_size=1)
        self.input_bn = nn.BatchNorm1d(channels[0])
        
        # Encoder blocks
        self.encoder_blocks = nn.ModuleList()
        
        for i in range(len(channels) - 1):
            block = self._make_encoder_block(
                channels[i], channels[i + 1], dropout
            )
            self.encoder_blocks.append(block)
        
        # Global context
        self.global_context = nn.Sequential(
            nn.Conv1d(channels[-1], channels[-1], kernel_size=1),
            nn.BatchNorm1d(channels[-1]),
            nn.ReLU(),
            nn.Dropout(dropout)
        )
        
        # Pooling
        self.pooling = AttentiveStatisticsPooling(channels[-1])
        
        # Embedding layer
        self.embedding_layer = nn.Sequential(
            nn.Linear(channels[-1] * 2, embedding_dim),
            nn.BatchNorm1d(embedding_dim),
            nn.ReLU(),
            nn.Dropout(dropout)
        )
        
        # Classification layer
        self.classifier = nn.Linear(embedding_dim, num_speakers)
        
        self._initialize_weights()
    
    def _make_encoder_block(self, in_channels: int, out_channels: int, dropout: float):
        """Create encoder block."""
        return nn.Sequential(
            # Depthwise separable convolution
            nn.Conv1d(in_channels, in_channels, kernel_size=3, padding=1, groups=in_channels),
            nn.Conv1d(in_channels, out_channels, kernel_size=1),
            nn.BatchNorm1d(out_channels),
            nn.ReLU(),
            nn.Dropout(dropout),
            
            # Pointwise convolution
            nn.Conv1d(out_channels, out_channels, kernel_size=1),
            nn.BatchNorm1d(out_channels),
            nn.ReLU(),
            nn.Dropout(dropout)
        )
    
    def _initialize_weights(self):
        """Initialize model weights."""
        for m in self.modules():
            if isinstance(m, nn.Conv1d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.xavier_normal_(m.weight)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm1d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
    
    def forward(self, x: torch.Tensor, extract_embedding: bool = False) -> torch.Tensor:
        """Forward pass through TitaNet-S."""
        # Input projection
        x = self.input_conv(x)
        x = self.input_bn(x)
        x = F.relu(x)
        
        # Encoder blocks
        for block in self.encoder_blocks:
            x = block(x)
        
        # Global context
        x = self.global_context(x)
        
        # Pooling
        x = self.pooling(x)
        
        # Embedding
        embeddings = self.embedding_layer(x)
        
        if extract_embedding:
            return F.normalize(embeddings, p=2, dim=1)
        
        # Classification
        logits = self.classifier(embeddings)
        return logits
    
    def extract_embedding(self, x: torch.Tensor) -> torch.Tensor:
        """Extract speaker embeddings."""
        return self.forward(x, extract_embedding=True)


class SpeakerNetMModel(nn.Module):
    """
    SpeakerNet-M (Medium) architecture.
    
    LSTM-based architecture with attention mechanism for speaker recognition.
    Good balance between performance and computational efficiency.
    """
    
    def __init__(self,
                 input_dim: int = 40,
                 embedding_dim: int = 256,
                 hidden_dim: int = 512,
                 num_layers: int = 4,
                 num_speakers: int = 4,
                 dropout: float = 0.1,
                 bidirectional: bool = True):
        super().__init__()
        
        self.input_dim = input_dim
        self.embedding_dim = embedding_dim
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.num_speakers = num_speakers
        self.bidirectional = bidirectional
        
        # Input projection
        self.input_proj = nn.Linear(input_dim, hidden_dim)
        
        # LSTM layers
        self.lstm = nn.LSTM(
            input_size=hidden_dim,
            hidden_size=hidden_dim,
            num_layers=num_layers,
            dropout=dropout if num_layers > 1 else 0,
            bidirectional=bidirectional,
            batch_first=True
        )
        
        lstm_output_dim = hidden_dim * (2 if bidirectional else 1)
        
        # Attention mechanism
        self.attention = nn.Sequential(
            nn.Linear(lstm_output_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, 1)
        )
        
        # Embedding layer
        self.embedding_layer = nn.Sequential(
            nn.Linear(lstm_output_dim, embedding_dim),
            nn.BatchNorm1d(embedding_dim),
            nn.ReLU(),
            nn.Dropout(dropout)
        )
        
        # Classification layer
        self.classifier = nn.Linear(embedding_dim, num_speakers)
        
        self._initialize_weights()
    
    def _initialize_weights(self):
        """Initialize model weights."""
        for name, param in self.named_parameters():
            if 'weight_ih' in name:
                nn.init.xavier_normal_(param.data)
            elif 'weight_hh' in name:
                nn.init.orthogonal_(param.data)
            elif 'bias' in name:
                nn.init.constant_(param.data, 0)
        
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_normal_(m.weight)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm1d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
    
    def forward(self, x: torch.Tensor, extract_embedding: bool = False) -> torch.Tensor:
        """
        Forward pass through SpeakerNet-M.
        
        Args:
            x: Input features [batch_size, input_dim, time]
            extract_embedding: Whether to return embeddings
            
        Returns:
            Either embeddings or classification logits
        """
        # Transpose to [batch_size, time, input_dim] for LSTM
        x = x.transpose(1, 2)
        batch_size, seq_len, _ = x.size()
        
        # Input projection
        x = self.input_proj(x)
        x = F.relu(x)
        
        # LSTM processing
        lstm_out, _ = self.lstm(x)  # [batch_size, seq_len, hidden_dim * directions]
        
        # Attention-based pooling
        attention_weights = self.attention(lstm_out)  # [batch_size, seq_len, 1]
        attention_weights = F.softmax(attention_weights, dim=1)
        
        # Weighted sum
        weighted_output = torch.sum(attention_weights * lstm_out, dim=1)  # [batch_size, hidden_dim * directions]
        
        # Embedding extraction
        embeddings = self.embedding_layer(weighted_output)
        
        if extract_embedding:
            return F.normalize(embeddings, p=2, dim=1)
        
        # Classification
        logits = self.classifier(embeddings)
        return logits
    
    def extract_embedding(self, x: torch.Tensor) -> torch.Tensor:
        """Extract speaker embeddings."""
        return self.forward(x, extract_embedding=True)


def create_model_from_config(architecture_type: Union[str, ModelArchitectureType],
                           input_dim: int,
                           embedding_dim: int,
                           num_speakers: int,
                           **kwargs) -> nn.Module:
    """
    Factory function to create model from configuration.
    
    Args:
        architecture_type: Type of architecture
        input_dim: Input feature dimension
        embedding_dim: Output embedding dimension
        num_speakers: Number of speakers
        **kwargs: Additional architecture-specific parameters
        
    Returns:
        Configured model
    """
    if isinstance(architecture_type, str):
        architecture_type = ModelArchitectureType(architecture_type)
    
    if architecture_type == ModelArchitectureType.ECAPA_TDNN:
        return ECADATDNNModel(
            input_dim=input_dim,
            embedding_dim=embedding_dim,
            num_speakers=num_speakers,
            channels=kwargs.get('ecapa_channels', 512),
            dropout=kwargs.get('dropout_rate', 0.1),
            use_attention_pooling=kwargs.get('use_attention_pooling', True)
        )
    
    elif architecture_type == ModelArchitectureType.TITANET_S:
        return TitaNetSModel(
            input_dim=input_dim,
            embedding_dim=embedding_dim,
            num_speakers=num_speakers,
            channels=kwargs.get('titanet_channels', [128, 256, 512, 1024]),
            dropout=kwargs.get('dropout_rate', 0.05)
        )
    
    elif architecture_type == ModelArchitectureType.SPEAKERNET_M:
        return SpeakerNetMModel(
            input_dim=input_dim,
            embedding_dim=embedding_dim,
            num_speakers=num_speakers,
            hidden_dim=kwargs.get('speakernet_hidden_dim', 512),
            num_layers=kwargs.get('speakernet_num_layers', 4),
            dropout=kwargs.get('dropout_rate', 0.1)
        )
    
    else:
        raise ValueError(f"Unsupported architecture type: {architecture_type}")


def get_supported_architectures() -> List[ModelArchitectureType]:
    """Get list of supported model architectures."""
    return list(ModelArchitectureType)


def count_parameters(model: nn.Module) -> Dict[str, int]:
    """Count the number of parameters in a model."""
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total_params = sum(p.numel() for p in model.parameters())
    
    return {
        'trainable_parameters': trainable_params,
        'total_parameters': total_params,
        'non_trainable_parameters': total_params - trainable_params
    }


def estimate_model_size(model: nn.Module) -> Dict[str, float]:
    """Estimate model size and memory requirements."""
    param_count = count_parameters(model)
    
    # Estimate size in MB (assuming float32)
    param_size_mb = param_count['total_parameters'] * 4 / (1024 * 1024)
    
    # Estimate memory for training (rough approximation)
    # Parameters + gradients + optimizer states (Adam: 2x params) + activations
    training_memory_mb = param_size_mb * 4  # Conservative estimate
    
    # Estimate memory for inference
    inference_memory_mb = param_size_mb * 1.5  # Parameters + activations
    
    return {
        'model_size_mb': param_size_mb,
        'inference_memory_mb': inference_memory_mb,
        'training_memory_mb': training_memory_mb
    }


# Example usage and testing
def test_model_architectures():
    """Test all model architectures with sample input."""
    batch_size = 4
    input_dim = 40
    seq_len = 100
    embedding_dim = 192
    num_speakers = 8
    
    # Create sample input
    x = torch.randn(batch_size, input_dim, seq_len)
    
    architectures = [
        (ModelArchitectureType.ECAPA_TDNN, {}),
        (ModelArchitectureType.TITANET_S, {}),
        (ModelArchitectureType.SPEAKERNET_M, {})
    ]
    
    for arch_type, kwargs in architectures:
        print(f"\nTesting {arch_type.value}:")
        
        # Create model
        model = create_model_from_config(
            arch_type, input_dim, embedding_dim, num_speakers, **kwargs
        )
        
        # Test forward pass
        model.eval()
        with torch.no_grad():
            # Classification output
            logits = model(x)
            print(f"  Logits shape: {logits.shape}")
            
            # Embedding extraction
            embeddings = model.extract_embedding(x)
            print(f"  Embeddings shape: {embeddings.shape}")
        
        # Model statistics
        param_stats = count_parameters(model)
        size_stats = estimate_model_size(model)
        
        print(f"  Parameters: {param_stats['trainable_parameters']:,}")
        print(f"  Model size: {size_stats['model_size_mb']:.2f} MB")
        print(f"  Inference memory: {size_stats['inference_memory_mb']:.2f} MB")


if __name__ == "__main__":
    test_model_architectures()