SCBs.py

import torch
from torch import nn
from transformers import WhisperConfig
from transformers.activations import ACT2FN
from transformers.models.whisper.modeling_whisper import WHISPER_ATTENTION_CLASSES
import torch.nn.functional as F
from .coattention import CoAttention
from .layers import CustomLinear, CustomDiagonalLinear, Gate

class LowRankApproxSelectFirst(nn.Module):
    def __init__(self, d_in, d_out, rank):
        super().__init__()
        self.d_in = d_in
        self.d_out = d_out
        self.rank = rank
        self.proj_in = nn.Linear(d_in, rank)
        self.proj_out = nn.Linear(rank, d_out)

    def forward(self, x):
        return self.proj_out(self.proj_in(x))

    def _init_weights(self):
        # Create low-rank approximation of the identity projection from first d_out of input
        eye = torch.eye(self.d_out, self.d_in)  # (d_out x d_in)

        # Low-rank SVD of eye matrix
        U, S, Vh = torch.linalg.svd(eye, full_matrices=False)  # U: (d_out x d_out), Vh: (d_in x d_in)

        U_k = U[:, :self.rank]              # (d_out x rank)
        S_k = S[:self.rank]                 # (rank,)
        V_k = Vh[:self.rank, :]             # (rank x d_in)

        A = V_k                             # (rank x d_in)
        B = U_k @ torch.diag(S_k)           # (d_out x rank)

        # Set weights
        self.proj_in.weight.data.copy_(A)
        self.proj_in.bias.data.zero_()
        self.proj_out.weight.data.copy_(B)
        self.proj_out.bias.data.zero_()



class TACBlock(nn.Module):
    def __init__(self, config: WhisperConfig, d_int_factor: float = 1, num_speakers=2):
        super().__init__()
        d = config.d_model
        d_prime = int(d * d_int_factor)
        self.num_speakers = num_speakers
        self.proj_in_1 = nn.Linear(d, d_prime, bias=True)
        self.proj_in_2 = nn.Linear(d, d_prime, bias=True)
        self.proj_int = nn.Linear(d_prime, d_prime,bias=True)
        self.proj_out_1 = nn.Linear(d+d_prime, d,bias=True)
        self.proj_out_2 = nn.Linear(d+d_prime, d,bias=True)
        self.activation_fn = ACT2FN[config.activation_function]
        self.norms = nn.ModuleList([nn.LayerNorm(d) for _ in range(self.num_speakers)])
        self.gate = Gate(self.num_speakers, 0.01)

    def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
        # hidden_states: (B, self.num_speakers, T, F)

        x_proj = torch.stack([self.activation_fn(self.proj_in_1(hidden_states[:,0])), self.activation_fn(self.proj_in_2(hidden_states[:, 1]))], dim=1)  # (B, 2, T, d')
        x_mean = x_proj.mean(dim=1, keepdim=True)  # (B, 1, T, d')
        z = self.activation_fn(self.proj_int(x_mean))  # (B, 1, T, d')

        z_expand = z.expand(-1, self.num_speakers, -1, -1)  # (B, self.num_speakers, T, d')
        x_cat = torch.cat([hidden_states, z_expand], dim=-1)  # (B, self.num_speakers, T, d + d')
        x_out = torch.stack([self.norms[0](self.proj_out_1(x_cat[:, 0])), self.norms[1](self.proj_out_2(x_cat[:, 1]))], dim=1)  # (B, self.num_speakers, T, d)
        return hidden_states + self.gate(x_out, dim=1)


class CrossAttentionBlock(nn.Module):
    def __init__(self, config: WhisperConfig):
        super().__init__()
        self.embed_dim = config.d_model

        self.num_speakers = getattr(config, "mt_num_speakers", 2)
        if self.num_speakers != 2:
            raise ValueError("CrossAttentionBlock supports only 2 speakers.")

        # Separate attention block per speaker
        self.attn_blocks = nn.ModuleList([
            WHISPER_ATTENTION_CLASSES[config._attn_implementation](
                embed_dim=self.embed_dim,
                num_heads=config.encoder_attention_heads,
                dropout=config.attention_dropout,
                config=config,
            )
            for _ in range(self.num_speakers)
        ])

        self.norms = nn.ModuleList([nn.LayerNorm(self.embed_dim) for _ in range(self.num_speakers)])
        self.gate = Gate(self.num_speakers, 0.01)

    def forward(self, hidden_states):
        # hidden_states: (B, 2, T, F)
        outputs = []
        for s in range(self.num_speakers):
            q = hidden_states[:, s]  # (B, T, F)
            other_s = 1 - s
            kv = hidden_states[:, other_s]  # (B, T, F)

            attn_out, _, _ = self.attn_blocks[s](hidden_states=q, key_value_states=kv)  # (B, T, F)
            outputs.append(self.norms[s](attn_out[:, None, :, :]))
        outputs =  torch.concat(outputs, dim=1)
        outputs_modulated = self.gate(outputs, dim=1) + hidden_states
        return outputs_modulated


class CompetitiveCrossAttentionBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.embed_dim = config.d_model
        self.num_heads = config.encoder_attention_heads
        self.head_dim = self.embed_dim // self.num_heads
        assert (
                self.head_dim * self.num_heads == self.embed_dim
        ), "embed_dim must be divisible by num_heads"

        self.num_speakers = getattr(config, "mt_num_speakers", 2)
        if self.num_speakers != 2:
            raise ValueError("CompetitiveCrossAttentionBlock supports only 2 speakers.")

        # Separate projections for Q, K, V
        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)

        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)

        self.norms = nn.ModuleList([nn.LayerNorm(self.embed_dim) for _ in range(self.num_speakers)])
        self.eps = 1e-6
        self.gate = Gate(self.num_speakers, 0.01)

    def _shape(self, tensor, seq_len, batch_size):
        # reshape into (B, num_heads, T, head_dim)
        return tensor.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

    def forward(self, hidden_states):
        # hidden_states: (B, 2, T, F)
        B, _, T, _ = hidden_states.shape

        h1, h2 = hidden_states[:, 0], hidden_states[:, 1]  # (B, T, F)

        # Project Q,K,V
        Q1 = self.q_proj(h1)  # (B, T, F)
        K2 = self.k_proj(h2)
        V2 = self.v_proj(h2)

        Q2 = self.q_proj(h2)
        K1 = self.k_proj(h1)
        V1 = self.v_proj(h1)

        # Reshape for multi-head attention
        Q1 = self._shape(Q1, T, B)  # (B, heads, T, head_dim)
        K2 = self._shape(K2, T, B)
        V2 = self._shape(V2, T, B)

        Q2 = self._shape(Q2, T, B)
        K1 = self._shape(K1, T, B)
        V1 = self._shape(V1, T, B)

        # Scaled dot-product attention logits
        scale = 1 / (self.head_dim ** 0.5)
        L_1to2 = torch.matmul(Q1, K2.transpose(-1, -2)) * scale  # (B, heads, T, T)
        L_2to1 = torch.matmul(Q2, K1.transpose(-1, -2)) * scale  # (B, heads, T, T)

        # Softmax over last dim (keys)
        S_1to2 = F.softmax(L_1to2, dim=-1)
        S_2to1 = F.softmax(L_2to1, dim=-1)

        # Competitive normalization (soft exclusivity)
        M_joint = S_1to2 + S_2to1 + self.eps
        A_1to2 = S_1to2 / M_joint
        A_2to1 = S_2to1 / M_joint

        # Weighted sum of values
        H1_attn = torch.matmul(A_1to2, V2)  # (B, heads, T, head_dim)
        H2_attn = torch.matmul(A_2to1, V1)

        # Concatenate heads back
        H1_attn = H1_attn.transpose(1, 2).contiguous().view(B, T, self.embed_dim)  # (B, T, F)
        H2_attn = H2_attn.transpose(1, 2).contiguous().view(B, T, self.embed_dim)

        # Output projection
        H1_attn = self.norms[0](self.out_proj(H1_attn))
        H2_attn = self.norms[1](self.out_proj(H2_attn))

        # Residuals
        out = hidden_states + self.gate(torch.concat([H1_attn[:, None, :, :], H2_attn[:, None, :, :]], dim=1), dim=1)

        return out # (B, 2, T, F)


class CoAttentionWrapper(nn.Module):
    def __init__(self, config, num_speakers=2):
        super().__init__()
        self.coa = CoAttention(embed_dim=config.d_model, single_dim=config.d_model//2, multi_dim=config.d_model // 4, n_heads=config.encoder_attention_heads, attn_dropout=config.attention_dropout)
        self.gate = Gate(num_speakers, 0.01)

    def forward(self, coa_input: torch.Tensor) -> torch.Tensor:
        # hidden_states: (B, 2, T, F)
        hidden_states = coa_input.permute(-2, 0, 1, -1)
        hidden_states = self.coa(hidden_states)
        out = coa_input + self.gate(hidden_states.permute(1, 2, 0, -1), dim=1)
        return out


class SpeakerCommunicationBlock(nn.Module):
    def __init__(self, config, scb_method):
        super().__init__()
        self.num_speakers = getattr(config, "mt_num_speakers", 2)
        self.embed_dim = config.d_model
        self.scb_method = scb_method
        self.config = config

        if self.scb_method == "tac":
            self.method = TACBlock(config)
        elif self.scb_method == "cross_attention":
            self.method = CrossAttentionBlock(config)
        elif self.scb_method == "competitive_cross_attention":
            self.method = CompetitiveCrossAttentionBlock(config)
        elif self.scb_method == "co_attention":
            self.method = CoAttentionWrapper(config)
        elif self.scb_method == "identity":
            self.method = (nn.Parameter(torch.zeros(self.embed_dim)) if config.fddt_bias_only else (
                CustomDiagonalLinear(self.embed_dim, bias=True, init_eye_val=1.0) if config.fddt_is_diagonal else CustomLinear(
                    self.embed_dim, self.embed_dim, bias=True, init_eye_val=1.0)))
        else:
            raise ValueError(f"Unsupported scb_method: {self.scb_method}")

    def forward(self, x):
        # x: (B, T, F)
        B, T, F = x.shape
        S = self.num_speakers

        # Reshape to (B//S, S, T, F)
        x_reshaped = x.view(B//S, S, T, F)

        # Call the selected method
        out = self.method(x_reshaped)

        # Reshape back (B, T, F)
        out_merged = out.view(B, T, F)
        return out_merged