| | |
|---|
| | | |
|---|
| | | import torch |
|---|
| | | import numpy as np |
|---|
| | | from typeguard import check_argument_types |
|---|
| | | |
|---|
| | | from funasr.models.encoder.abs_encoder import AbsEncoder |
|---|
| | | from funasr.models.frontend.abs_frontend import AbsFrontend |
|---|
| | |
|---|
| | | from funasr.modules.add_sos_eos import add_sos_eos |
|---|
| | | from funasr.modules.nets_utils import make_pad_mask, pad_list |
|---|
| | | from funasr.torch_utils.device_funcs import force_gatherable |
|---|
| | | from funasr.train.abs_espnet_model import AbsESPnetModel |
|---|
| | | from funasr.models.base_model import FunASRModel |
|---|
| | | from funasr.models.predictor.cif import CifPredictorV3 |
|---|
| | | |
|---|
| | | |
|---|
| | | if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): |
|---|
| | | from torch.cuda.amp import autocast |
|---|
| | |
|---|
| | | yield |
|---|
| | | |
|---|
| | | |
|---|
| | | class TimestampPredictor(AbsESPnetModel): |
|---|
| | | class TimestampPredictor(FunASRModel): |
|---|
| | | """ |
|---|
| | | Author: Speech Lab, Alibaba Group, China |
|---|
| | | Author: Speech Lab of DAMO Academy, Alibaba Group |
|---|
| | | """ |
|---|
| | | |
|---|
| | | def __init__( |
|---|
| | |
|---|
| | | predictor_bias: int = 0, |
|---|
| | | token_list=None, |
|---|
| | | ): |
|---|
| | | assert check_argument_types() |
|---|
| | | |
|---|
| | | super().__init__() |
|---|
| | | # note that eos is the same as sos (equivalent ID) |
|---|
| | |
|---|
| | | self.predictor_bias = predictor_bias |
|---|
| | | self.criterion_pre = mae_loss() |
|---|
| | | self.token_list = token_list |
|---|
| | | |
|---|
| | | |
|---|
| | | def forward( |
|---|
| | | self, |
|---|
| | | speech: torch.Tensor, |
|---|
| | |
|---|
| | | text_lengths: torch.Tensor, |
|---|
| | | ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]: |
|---|
| | | """Frontend + Encoder + Decoder + Calc loss |
|---|
| | | |
|---|
| | | Args: |
|---|
| | | speech: (Batch, Length, ...) |
|---|
| | | speech_lengths: (Batch, ) |
|---|
| | |
|---|
| | | self, speech: torch.Tensor, speech_lengths: torch.Tensor |
|---|
| | | ) -> Tuple[torch.Tensor, torch.Tensor]: |
|---|
| | | """Frontend + Encoder. Note that this method is used by asr_inference.py |
|---|
| | | |
|---|
| | | Args: |
|---|
| | | speech: (Batch, Length, ...) |
|---|
| | | speech_lengths: (Batch, ) |
|---|
| | |
|---|
| | | encoder_out, encoder_out_lens, _ = self.encoder(feats, feats_lengths) |
|---|
| | | |
|---|
| | | return encoder_out, encoder_out_lens |
|---|
| | | |
|---|
| | | |
|---|
| | | def _extract_feats( |
|---|
| | | self, speech: torch.Tensor, speech_lengths: torch.Tensor |
|---|
| | | ) -> Tuple[torch.Tensor, torch.Tensor]: |
|---|
| | |
|---|
| | | def calc_predictor_timestamp(self, encoder_out, encoder_out_lens, token_num): |
|---|
| | | encoder_out_mask = (~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :]).to( |
|---|
| | | encoder_out.device) |
|---|
| | | ds_alphas, ds_cif_peak, us_alphas, us_cif_peak = self.predictor.get_upsample_timestamp(encoder_out, |
|---|
| | | encoder_out_mask, |
|---|
| | | token_num) |
|---|
| | | return ds_alphas, ds_cif_peak, us_alphas, us_cif_peak |
|---|
| | | ds_alphas, ds_cif_peak, us_alphas, us_peaks = self.predictor.get_upsample_timestamp(encoder_out, |
|---|
| | | encoder_out_mask, |
|---|
| | | token_num) |
|---|
| | | return ds_alphas, ds_cif_peak, us_alphas, us_peaks |
|---|
| | | |
|---|
| | | def collect_feats( |
|---|
| | | self, |
|---|
