Export forward. Args: speech: Speech audio tensor, shape (batch, time). speech_lengths: Length of each speech sample.
(
self,
speech: torch.Tensor,
speech_lengths: torch.Tensor,
)
| 42 | |
| 43 | |
| 44 | def export_forward( |
| 45 | self, |
| 46 | speech: torch.Tensor, |
| 47 | speech_lengths: torch.Tensor, |
| 48 | ): |
| 49 | # a. To device |
| 50 | """Export forward. |
| 51 | |
| 52 | Args: |
| 53 | speech: Speech audio tensor, shape (batch, time). |
| 54 | speech_lengths: Length of each speech sample. |
| 55 | """ |
| 56 | batch = {"speech": speech, "speech_lengths": speech_lengths} |
| 57 | |
| 58 | enc, enc_len = self.encoder(**batch) |
| 59 | mask = self.make_pad_mask(enc_len)[:, None, :] |
| 60 | pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index = self.predictor(enc, mask) |
| 61 | pre_token_length = pre_token_length.round().type(torch.int32) |
| 62 | |
| 63 | decoder_out, _ = self.decoder(enc, enc_len, pre_acoustic_embeds, pre_token_length) |
| 64 | decoder_out = torch.log_softmax(decoder_out, dim=-1) |
| 65 | |
| 66 | # get predicted timestamps |
| 67 | us_alphas, us_cif_peak = self.predictor.get_upsample_timestmap(enc, mask, pre_token_length) |
| 68 | |
| 69 | return decoder_out, pre_token_length, us_alphas, us_cif_peak |
| 70 | |
| 71 | |
| 72 | def export_dummy_inputs(self): |
nothing calls this directly
no test coverage detected
searching dependent graphs…