| 194 | |
| 195 | class WhisperEncoder(nn.Module): |
| 196 | def __init__( |
| 197 | self, |
| 198 | in_dim: int, |
| 199 | embed_dim: int, |
| 200 | num_layers: int, |
| 201 | num_heads: int, |
| 202 | ffn_dim: int = None, |
| 203 | attn_dropout: float = 0.0, |
| 204 | dropout: float = 0.0, |
| 205 | max_positions: int = 1500, |
| 206 | ): |
| 207 | super().__init__() |
| 208 | self.in_dim = in_dim |
| 209 | self.embed_dim = embed_dim |
| 210 | self.dropout = dropout |
| 211 | # Input downsampling |
| 212 | self.conv1 = nn.Conv1d(in_dim, embed_dim, kernel_size=3, padding=1) |
| 213 | self.conv2 = nn.Conv1d(embed_dim, embed_dim, kernel_size=3, stride=2, padding=1) |
| 214 | # Fixed positional embedding |
| 215 | self.max_positions = max_positions |
| 216 | self.embed_positions = nn.Embedding(self.max_positions, embed_dim) |
| 217 | self.embed_positions.requires_grad_(False) |
| 218 | # Transformer |
| 219 | self.layers = nn.ModuleList( |
| 220 | [ |
| 221 | WhisperEncoderLayer( |
| 222 | embed_dim, num_heads, ffn_dim, attn_dropout, dropout |
| 223 | ) |
| 224 | for _ in range(num_layers) |
| 225 | ] |
| 226 | ) |
| 227 | # Output norm |
| 228 | self.layer_norm = nn.LayerNorm(embed_dim) |
| 229 | # Init weight |
| 230 | self.apply(self._init_weights) |
| 231 | # Init position embedding |
| 232 | self.embed_positions.weight.copy_(sinusoids(*self.embed_positions.weight.shape)) |
| 233 | |
| 234 | def forward( |
| 235 | self, |