Class WhisperEncoder

fireredtts2/codec/whisper.py:195–272

class WhisperEncoder(nn.Module):
    def __init__(
        self,
        in_dim: int,
        embed_dim: int,
        num_layers: int,
        num_heads: int,
        ffn_dim: int = None,
        attn_dropout: float = 0.0,
        dropout: float = 0.0,
        max_positions: int = 1500,
    ):
        super().__init__()
        self.in_dim = in_dim
        self.embed_dim = embed_dim
        self.dropout = dropout
        # Input downsampling
        self.conv1 = nn.Conv1d(in_dim, embed_dim, kernel_size=3, padding=1)
        self.conv2 = nn.Conv1d(embed_dim, embed_dim, kernel_size=3, stride=2, padding=1)
        # Fixed positional embedding
        self.max_positions = max_positions
        self.embed_positions = nn.Embedding(self.max_positions, embed_dim)
        self.embed_positions.requires_grad_(False)
        # Transformer
        self.layers = nn.ModuleList(
            [
                WhisperEncoderLayer(
                    embed_dim, num_heads, ffn_dim, attn_dropout, dropout
                )
                for _ in range(num_layers)
            ]
        )
        # Output norm
        self.layer_norm = nn.LayerNorm(embed_dim)
        # Init weight
        self.apply(self._init_weights)
        # Init position embedding
        self.embed_positions.weight.copy_(sinusoids(*self.embed_positions.weight.shape))

    def forward(
        self,
        hidden_states: torch.Tensor,
        hidden_length: torch.Tensor,
        apply_position: bool = True,
    ):
        # Downsampling
        hidden_states = hidden_states.transpose(1, 2)
        hidden_states = F.gelu(self.conv1(hidden_states))
        hidden_states = F.gelu(self.conv2(hidden_states))
        hidden_states = hidden_states.transpose(1, 2)
        hidden_length = hidden_length // 2  # from 100Hz -> 50Hz
        # Pos encoding
        if apply_position:
            pos_embed = self.embed_positions(
                torch.arange(0, hidden_states.shape[1], device=hidden_states.device)
            )
            hidden_states = hidden_states + pos_embed
        hidden_states = nn.functional.dropout(