Spaces:
Sleeping
Sleeping
| import copy | |
| from typing import List, Optional, Tuple, Union | |
| import torch | |
| from torch import Tensor, nn | |
| from .encoder_layer import BypassModule | |
| from .scaling import FloatLike | |
| class DownsampledZipformer2Encoder(nn.Module): | |
| r""" | |
| DownsampledZipformer2Encoder is a zipformer encoder evaluated at a reduced frame rate, | |
| after convolutional downsampling, and then upsampled again at the output, and combined | |
| with the origin input, so that the output has the same shape as the input. | |
| """ | |
| def __init__(self, encoder: nn.Module, dim: int, downsample: int, dropout: FloatLike) -> None: | |
| super(DownsampledZipformer2Encoder, self).__init__() | |
| self.downsample_factor = downsample | |
| self.downsample = SimpleDownsample(dim, downsample, dropout) | |
| self.num_layers = encoder.num_layers | |
| self.encoder = encoder | |
| self.upsample = SimpleUpsample(dim, downsample) | |
| self.out_combiner = BypassModule(dim, straight_through_rate=0) | |
| def forward( | |
| self, | |
| src: Tensor, | |
| chunk_size: int = -1, | |
| feature_mask: Union[Tensor, float] = 1.0, | |
| attn_mask: Optional[Tensor] = None, | |
| src_key_padding_mask: Optional[Tensor] = None, | |
| ) -> Tensor: | |
| r"""Downsample, go through encoder, upsample. | |
| Args: | |
| src: the sequence to the encoder (required): shape (seq_len, batch_size, embedding_dim). | |
| feature_mask: something that broadcasts with src, that we'll multiply `src` | |
| by at every layer: if a Tensor, likely of shape (seq_len, batch_size, embedding_dim) | |
| attn_mask: the attention mask, of shape (batch_size, seq_len, seq_len) or (seq_len, seq_len), | |
| interpreted as (batch_size, tgt_seq_len, src_seq_len) or (tgt_seq_len, src_seq_len). | |
| True means masked position. May be None. | |
| src_key_padding_mask: the mask for padding, of shape (batch_size, seq_len); True means | |
| masked position. May be None. | |
| Returns: a Tensor with the same shape as src. | |
| """ | |
| src_orig = src | |
| ds = self.downsample_factor | |
| src = self.downsample(src, src_key_padding_mask) | |
| if attn_mask is not None: | |
| attn_mask = attn_mask[::ds, ::ds] | |
| downsampled_padding_mask = None if src_key_padding_mask is None else src_key_padding_mask[..., ::ds] | |
| src = self.encoder( | |
| src, | |
| chunk_size=chunk_size // ds, | |
| feature_mask=feature_mask, | |
| attn_mask=attn_mask, | |
| src_key_padding_mask=downsampled_padding_mask, | |
| ) | |
| src = self.upsample(src) | |
| # remove any extra frames that are not a multiple of downsample_factor | |
| src = src[: src_orig.shape[0]] | |
| return self.out_combiner(src_orig, src) | |
| def streaming_forward( | |
| self, | |
| src: Tensor, | |
| states: List[Tensor], | |
| left_context_len: int, | |
| src_key_padding_mask: Tensor, | |
| ) -> Tuple[Tensor, List[Tensor]]: | |
| r"""Downsample, go through encoder, upsample, in streaming forward mode. | |
| Args: | |
| src: the sequence to the encoder (required): shape (seq_len, batch_size, embedding_dim). | |
| states: list of cached tensors of N encoder layers. For layer-i, states[i*6:(i+1)*6] is | |
| (cached_key, cached_nonlin_attn, cached_val1, cached_val2, cached_conv1, cached_conv2). | |
| left_context_len: Number of left context frames. | |
| src_key_padding_mask: the mask for padding, of shape (batch_size, left_context_len+seq_len); | |
| True means masked position. May be None. | |
| Returns: | |
| - output, a Tensor with the same shape as src. | |
| - updated states | |
| """ | |
| src_orig = src | |
| src = self.downsample(src) | |
| src, new_states = self.encoder.streaming_forward( | |
| src, | |
| states=states, | |
| left_context_len=left_context_len, | |
| src_key_padding_mask=src_key_padding_mask, | |
| ) | |
| src = self.upsample(src) | |
| # remove any extra frames that are not a multiple of downsample_factor | |
| src = src[: src_orig.shape[0]] | |
| return self.out_combiner(src_orig, src), new_states | |
| class SimpleDownsample(torch.nn.Module): | |
| """ | |
| Does downsampling with attention, by weighted sum, and a projection.. | |
| """ | |
| def __init__( | |
| self, | |
| channels: int, | |
| downsample: int, | |
| dropout: FloatLike, | |
| ) -> None: | |
| super(SimpleDownsample, self).__init__() | |
| self.bias = nn.Parameter(torch.zeros(downsample)) | |
| self.name = None # will be set from training code | |
| self.dropout = copy.deepcopy(dropout) | |
| self.downsample = downsample | |
| def forward( | |
| self, | |
| src: Tensor, | |
| src_key_padding_mask: Optional[Tensor] = None, | |
| ) -> Tensor: | |
| """ | |
| x: (seq_len, batch_size, in_channels) | |
| Returns a tensor of shape | |
| ( (seq_len+downsample-1)//downsample, batch_size, channels) | |
| """ | |
| (seq_len, batch_size, in_channels) = src.shape | |
| ds = self.downsample | |
| d_seq_len = (seq_len + ds - 1) // ds | |
| # --------------------------------------------------------- | |
| # Match unpadded inference by extending each sequence's last valid frame. | |
| # --------------------------------------------------------- | |
| if src_key_padding_mask is not None: | |
| lengths = (~src_key_padding_mask).sum(dim=-1) | |
| last_valid_indices = lengths.sub(1).view(1, batch_size, 1) | |
| last_valid = src.gather( | |
| 0, | |
| last_valid_indices.expand(1, batch_size, in_channels), | |
| ) | |
| src = torch.where( | |
| src_key_padding_mask.transpose(0, 1).unsqueeze(-1), | |
| last_valid, | |
| src, | |
| ) | |
| # Pad to an exact multiple of self.downsample | |
| # right-pad src, repeating the last element. | |
| pad = d_seq_len * ds - seq_len | |
| src_extra = src[src.shape[0] - 1 :].expand(pad, src.shape[1], src.shape[2]) | |
| src = torch.cat((src, src_extra), dim=0) | |
| assert src.shape[0] == d_seq_len * ds | |
| src = src.reshape(d_seq_len, ds, batch_size, in_channels) | |
| weights = self.bias.softmax(dim=0) | |
| # weights: (downsample, 1, 1) | |
| weights = weights.unsqueeze(-1).unsqueeze(-1) | |
| # ans1 is the first `in_channels` channels of the output | |
| ans = (src * weights).sum(dim=1) | |
| return ans | |
| class SimpleUpsample(torch.nn.Module): | |
| """ | |
| A very simple form of upsampling that mostly just repeats the input, but | |
| also adds a position-specific bias. | |
| """ | |
| def __init__(self, num_channels: int, upsample: int) -> None: | |
| super(SimpleUpsample, self).__init__() | |
| self.upsample = upsample | |
| def forward(self, src: Tensor) -> Tensor: | |
| """ | |
| x: (seq_len, batch_size, num_channels) | |
| Returns a tensor of shape | |
| ( (seq_len*upsample), batch_size, num_channels) | |
| """ | |
| upsample = self.upsample | |
| (seq_len, batch_size, num_channels) = src.shape | |
| src = src.unsqueeze(1).expand(seq_len, upsample, batch_size, num_channels) | |
| src = src.reshape(seq_len * upsample, batch_size, num_channels) | |
| return src | |