Spaces:
Sleeping
Sleeping
| import torch | |
| from torch import nn | |
| class AudioAdapter(nn.Module): | |
| def __init__(self, input_size: int, hidden_size: int) -> None: | |
| super().__init__() | |
| self.input_norm = nn.LayerNorm(input_size) | |
| self.conv = nn.Conv1d( | |
| input_size, | |
| input_size, | |
| kernel_size=3, | |
| stride=2, | |
| padding=1, | |
| ) | |
| self.activation = nn.GELU() | |
| self.projection = nn.Linear(input_size, hidden_size) | |
| self.output_norm = nn.LayerNorm(hidden_size) | |
| def forward( | |
| self, | |
| hidden_states: torch.Tensor, | |
| padding_mask: torch.BoolTensor, | |
| ) -> tuple[torch.Tensor, torch.BoolTensor]: | |
| # --------------------------------------------------------- | |
| # Compress the audio sequence and change its hidden size. | |
| # --------------------------------------------------------- | |
| hidden_states = self.input_norm(hidden_states) | |
| hidden_states = hidden_states.masked_fill(padding_mask.unsqueeze(-1), 0.0) | |
| hidden_states = self.conv(hidden_states.transpose(1, 2)).transpose(1, 2) | |
| hidden_states = self.activation(hidden_states) | |
| hidden_states = self.projection(hidden_states) | |
| hidden_states = self.output_norm(hidden_states) | |
| # --------------------------------------------------------- | |
| # Derive exact convolution output lengths before batch padding. | |
| # --------------------------------------------------------- | |
| input_lengths = (~padding_mask).sum(dim=-1) | |
| output_lengths = (input_lengths + 1) // 2 | |
| positions = torch.arange(hidden_states.shape[1], device=hidden_states.device) | |
| output_padding_mask = positions.unsqueeze(0) >= output_lengths.unsqueeze(1) | |
| return hidden_states, output_padding_mask | |