| from typing import Tuple, Union |
| import torch |
| from wenet.transformer.subsampling import BaseSubsampling |
|
|
|
|
| class IdentitySubsampling(BaseSubsampling): |
| """ Paraformer subsampling |
| """ |
|
|
| def __init__(self, idim: int, odim: int, dropout_rate: float, |
| pos_enc_class: torch.nn.Module): |
| super().__init__() |
| _, _ = idim, odim |
| self.right_context = 6 |
| self.subsampling_rate = 6 |
| self.pos_enc = pos_enc_class |
|
|
| def forward( |
| self, |
| x: torch.Tensor, |
| x_mask: torch.Tensor, |
| offset: Union[torch.Tensor, int] = 0 |
| ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| """Subsample x. |
| |
| Args: |
| x (torch.Tensor): Input tensor (#batch, time, idim). |
| x_mask (torch.Tensor): Input mask (#batch, 1, time). |
| |
| Returns: |
| torch.Tensor: Subsampled tensor (#batch, time', odim), |
| where time' = time. |
| torch.Tensor: Subsampled mask (#batch, 1, time'), |
| where time' = time |
| torch.Tensor: positional encoding |
| |
| """ |
| |
| if isinstance(offset, torch.Tensor): |
| offset = torch.add(offset, 1) |
| else: |
| offset = offset + 1 |
| x, pos_emb = self.pos_enc(x, offset) |
| return x, pos_emb, x_mask |
|
|
| def position_encoding(self, offset: Union[int, torch.Tensor], |
| size: int) -> torch.Tensor: |
| return self.pos_enc.position_encoding(offset + 1, size) |
|
|
