speaker_embedding_architecture.py

from dataclasses import dataclass
from typing import Optional
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import Wav2Vec2Model, PreTrainedModel
from transformers.modeling_outputs import ModelOutput
from .ge2e_loss import GE2ELoss
from .speaker_embedding_config import SpeakerEmbeddingConfig


@dataclass
class EmbeddingOutput(ModelOutput):
    loss: Optional[torch.FloatTensor] = None
    embeddings: torch.Tensor = None  # shape (batch_size, embedding_size)
    dvecs: Optional[torch.Tensor] = None  # shape (N, M, embedding_size) for analysis


# https://arxiv.org/pdf/1803.10963
class AttentiveStatisticsPooling(nn.Module):
    def __init__(self, input_dim, hidden_dim=128):
        super().__init__()
        self.attention = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, 1)
        )

    def forward(self, x, mask=None):
        # x: (batch, time, features)

        # Compute attention weights
        attn_weights = self.attention(x)  # (batch, time, 1)

        if mask is not None:
            # mask: (batch, time), True = valid, False = padding
            attn_weights = attn_weights.masked_fill(~mask.unsqueeze(-1), float('-inf'))

        attn_weights = F.softmax(attn_weights, dim=1)  # (batch, time, 1)

        # Weighted mean
        mean = torch.sum(x * attn_weights, dim=1)  # (batch, features)

        # Weighted std
        variance = torch.sum(attn_weights * (x - mean.unsqueeze(1)) ** 2, dim=1)
        std = torch.sqrt(variance.clamp(min=1e-5))

        # Concatenate mean and std
        return torch.cat([mean, std], dim=-1)  # (batch, features * 2)


class ZeroModule(nn.Module):
    def forward(self, hidden_states):
        return torch.zeros_like(hidden_states)


class SpeakerEmbeddingArchitecture(PreTrainedModel):
    config_class = SpeakerEmbeddingConfig

    def __init__(self, config):
        super().__init__(config)
        self.config = config
        self.encoder = Wav2Vec2Model(config)
        self.encoder.init_weights()

        if getattr(config, 'disable_positional_embeddings', False):
            self.encoder.encoder.pos_conv_embed = ZeroModule()

        # Projection layer to get desired embedding dimension
        self.pooling = AttentiveStatisticsPooling(config.hidden_size, hidden_dim=config.hidden_size)

        self.embedding_size = getattr(config, 'embedding_size', 256)

        n_projection_layers = getattr(config, 'n_projection_layers', 1)
        self.projection_layers = nn.ModuleList([])
        for l in range(n_projection_layers - 1):
            self.projection_layers.append(
                nn.Sequential(
                    nn.Linear(
                        config.hidden_size * 2 if l == 0 else self.embedding_size,
                        self.embedding_size
                    ),
                    #nn.BatchNorm1d(self.embedding_size, momentum=0.01, eps=1e-5),
                    nn.LeakyReLU(0.01)
                )
            )

        self.final_projection = nn.Sequential(
            nn.Linear(
                config.hidden_size * 2 if n_projection_layers == 1 else self.embedding_size,
                self.embedding_size
            ),
            #nn.BatchNorm1d(self.embedding_size, momentum=0.01, eps=1e-5),
            nn.LeakyReLU(0.01)
        )

        self.layer_weights = None
        if getattr(config, 'use_layer_weights', False):
            num_layers = config.num_hidden_layers + 1
            self.layer_weights = nn.Parameter(torch.ones(num_layers), requires_grad=True)

        # Optional: Layer normalization before pooling
        self.layer_norm = nn.LayerNorm(config.hidden_size)
        self.margin = config.loss_margin
        self.scale = config.loss_scale

        for m in self.pooling.attention.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight, gain=0.2)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)

        # Normal init for projection layers
        self.projection_layers.apply(self._init_weights)
        self.final_projection.apply(self._init_weights)

        # GE2E Loss
        # self.ge2e_loss = GE2ELoss(loss_method='softmax')

    def _init_weights(self, module):
        """Initialize the weights"""
        if isinstance(module, nn.Linear):
            nn.init.xavier_uniform_(module.weight, gain=0.5)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.BatchNorm1d):
            nn.init.ones_(module.weight)
            nn.init.zeros_(module.bias)

    def temporal_pooling(self, hidden_states, attention_mask=None):
        """
        Pool over time dimension with optional attention mask.

        Args:
            hidden_states: (batch_size, seq_len, hidden_size)
            attention_mask: (batch_size, seq_len) - optional

        Returns:
            pooled: (batch_size, hidden_size)
        """
        if attention_mask is not None:
            # Mask padding tokens before pooling
            # attention_mask shape: (batch_size, seq_len)
            mask_expanded = attention_mask.unsqueeze(-1).expand(hidden_states.size())
            sum_hidden = torch.sum(hidden_states * mask_expanded, dim=1)
            sum_mask = torch.clamp(mask_expanded.sum(dim=1), min=1e-9)
            pooled = sum_hidden / sum_mask
        else:
            # Simple mean pooling
            pooled = torch.mean(hidden_states, dim=1)

        return pooled

    def forward(
            self,
            input_values: Optional[torch.Tensor],
            attention_mask: Optional[torch.Tensor] = None,
            labels: Optional[torch.Tensor] = None,
            mask_time_indices: Optional[torch.FloatTensor] = None,
            output_attentions: Optional[bool] = None,
            output_hidden_states: Optional[bool] = None,
            return_dict: Optional[bool] = None,
            cached_centroids: Optional[torch.Tensor] = None,
            cached_labels: Optional[torch.Tensor] = None,
    ):
        """
        Args:
            input_values: (N*M, seq_len) audio input
            attention_mask: (N*M, seq_len) optional
            compute_loss: whether to compute GE2E loss (needs N and M)
            speakers_per_batch: N - number of speakers in batch
            utterances_per_speaker: M - utterances per speaker

        Returns:
            EmbeddingOutput with loss and embeddings
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        # 1. Encode audio
        encoder_outputs = self.encoder(
            input_values=input_values,
            attention_mask=attention_mask,
            mask_time_indices=mask_time_indices,
            output_attentions=output_attentions,
            output_hidden_states=True,
            return_dict=return_dict,
        )

        if self.layer_weights is None:
            hidden_states = encoder_outputs.last_hidden_state
        else:
            # Weighted sum of hidden states from all layers
            hidden_states = torch.stack(encoder_outputs.hidden_states, dim=0)
            layer_weights = F.softmax(self.layer_weights, dim=0)
            hidden_states = (hidden_states * layer_weights.view(-1, 1, 1, 1)).sum(dim=0)

        # hidden_states = encoder_outputs.last_hidden_state  # (N*M, seq_len, hidden_size)

        sub_attention_mask = None
        if attention_mask is not None:
            # Use the built-in Wav2Vec2 function!
            sub_attention_mask = self.encoder._get_feature_vector_attention_mask(
                feature_vector_length=hidden_states.shape[1],
                attention_mask=attention_mask,
                add_adapter=False
            )

        # 2. Optional normalization
        hidden_states = self.layer_norm(hidden_states)

        # 3. Temporal pooling
        pooled = self.pooling(hidden_states, mask=sub_attention_mask)
        # pooled = self.temporal_pooling(hidden_states, sub_attention_mask)  # (N*M, hidden_size)

        # 4. Project to embedding space
        # embeddings = self.projection(pooled)  # (N*M, embedding_size)
        embeddings = pooled
        for projection_layer in self.projection_layers:
            embeddings = projection_layer(embeddings)

        embeddings = self.final_projection(embeddings)

        # 5. L2 normalize embeddings (important for GE2E and cosine similarity!)
        embeddings = nn.functional.normalize(embeddings, p=2, dim=1, eps=1e-8)

        if cached_centroids is not None:
            # We are inside the DDP Forward pass now!
            # DDP will see 'w' and 'b' being used here.

            # w = self.ge2e_loss.w
            # b = self.ge2e_loss.b

            # Calculate Similarity against CACHED centroids
            # embeddings: (Batch, D), cached_centroids: (N, D)
            # sim_matrix = torch.mm(embeddings, cached_centroids.transpose(0, 1))
            sim_matrix = torch.einsum('bd,bnd->bn', embeddings, cached_centroids)
            # sim_matrix = torch.clamp(sim_matrix, min=1e-6)
            # sim_matrix = w * sim_matrix + b

            # https://arxiv.org/pdf/1801.05599 this loss here
            one_hot = torch.zeros_like(sim_matrix)
            one_hot.scatter_(1, cached_labels.unsqueeze(1), 1.0)
            sim_matrix = sim_matrix - one_hot * self.margin
            sim_matrix = sim_matrix * self.scale

            if cached_labels is not None:
                # loss_fct = nn.CrossEntropyLoss()
                loss = F.cross_entropy(sim_matrix, cached_labels)
                # Return the Loss directly
                return EmbeddingOutput(loss=loss, embeddings=embeddings)
            else:
                return EmbeddingOutput(loss=None, embeddings=embeddings)

        return EmbeddingOutput(loss=None, embeddings=embeddings)

    def get_speaker_embedding(self, audio, attention_mask=None):
        """
        Convenience method to get embedding for a single audio sample.

        Args:
            audio: (1, seq_len) or (seq_len,)
            attention_mask: optional

        Returns:
            embedding: (embedding_size,)
        """
        if audio.dim() == 1:
            audio = audio.unsqueeze(0)

        output = self.forward(
            input_values=audio,
            attention_mask=attention_mask,
            compute_loss=False,
            return_dict=True
        )

        return output.embeddings.squeeze(0)