Spaces:
Sleeping
Sleeping
File size: 1,799 Bytes
82e11f0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 | 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
|