| 119 | |
| 120 | |
| 121 | class WhisperEncoderLayer(nn.Module): |
| 122 | def __init__( |
| 123 | self, |
| 124 | embed_dim: int, |
| 125 | num_heads: int, |
| 126 | ffn_dim: int = None, |
| 127 | attn_dropout: float = 0.0, |
| 128 | dropout: float = 0.0, |
| 129 | ): |
| 130 | super().__init__() |
| 131 | self.dropout = dropout |
| 132 | # Attention |
| 133 | self.self_attn = WhisperSdpaAttention(embed_dim, num_heads, attn_dropout) |
| 134 | self.self_attn_layer_norm = nn.LayerNorm(embed_dim) |
| 135 | # FFN |
| 136 | ffn_dim = ffn_dim if ffn_dim is not None else embed_dim * 4 |
| 137 | self.fc1 = nn.Linear(embed_dim, ffn_dim) |
| 138 | self.fc2 = nn.Linear(ffn_dim, embed_dim) |
| 139 | # Output norm |
| 140 | self.final_layer_norm = nn.LayerNorm(embed_dim) |
| 141 | |
| 142 | def forward( |
| 143 | self, |
| 144 | hidden_states: torch.Tensor, |
| 145 | attention_mask: torch.Tensor, |
| 146 | ): |
| 147 | # Attention |
| 148 | residual = hidden_states |
| 149 | hidden_states = self.self_attn_layer_norm(hidden_states) |
| 150 | hidden_states = self.self_attn(hidden_states, attention_mask) |
| 151 | hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training) |
| 152 | hidden_states = residual + hidden_states |
| 153 | |
| 154 | # FFN |
| 155 | residual = hidden_states |
| 156 | hidden_states = self.final_layer_norm(hidden_states) |
| 157 | hidden_states = F.gelu(self.fc1(hidden_states)) |
| 158 | hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training) |
| 159 | hidden_states = self.fc2(hidden_states) |
| 160 | hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training) |
| 161 | hidden_states = residual + hidden_states |
| 162 | return hidden_states |
| 163 | |
| 164 | def forward_chunk( |
| 165 | self, |
| 166 | hidden_states: torch.Tensor, |
| 167 | kv_cache: torch.Tensor = None, |
| 168 | ): |
| 169 | """Forward self-attention with kv cache. |
| 170 | |
| 171 | Args: |
| 172 | hidden_states: shape (b, t, c) |
| 173 | kv_cache: shape (b, nh, t, c*2) |
| 174 | """ |
| 175 | # Attention |
| 176 | residual = hidden_states |
| 177 | hidden_states = self.self_attn_layer_norm(hidden_states) |
| 178 | hidden_states, new_kv_cache = self.self_attn.forward_chunk( |