speechbrain/asr-wav2vec2-ctc-aishell

1"""Custom Interface for AISHELL-1 CTC inference2An external tokenizer is used so some special tokens3need to be specified during decoding4 5Authors6 * Yingzhi Wang 20227"""8 9import torch10from speechbrain.inference.interfaces import Pretrained11 12 13class CustomEncoderDecoderASR(Pretrained):14    """A ready-to-use Encoder-Decoder ASR model15    The class can be used either to run only the encoder (encode()) to extract16    features or to run the entire encoder-decoder model17    (transcribe()) to transcribe speech. The given YAML must contains the fields18    specified in the *_NEEDED[] lists.19    Example20    -------21    >>> from speechbrain.pretrained import EncoderDecoderASR22    >>> tmpdir = getfixture("tmpdir")23    >>> asr_model = EncoderDecoderASR.from_hparams(24    ...     source="speechbrain/asr-crdnn-rnnlm-librispeech",25    ...     savedir=tmpdir,26    ... )27    >>> asr_model.transcribe_file("tests/samples/single-mic/example2.flac")28    "MY FATHER HAS REVEALED THE CULPRIT'S NAME"29    """30 31    def __init__(self, *args, **kwargs):32        super().__init__(*args, **kwargs)33        self.tokenizer = self.hparams.tokenizer34 35    def transcribe_file(self, path):36        """Transcribes the given audiofile into a sequence of words.37        Arguments38        ---------39        path : str40            Path to audio file which to transcribe.41        Returns42        -------43        str44            The audiofile transcription produced by this ASR system.45        """46        waveform = self.load_audio(path)47        # Fake a batch:48        batch = waveform.unsqueeze(0)49        rel_length = torch.tensor([1.0])50        predicted_words = self.transcribe_batch(51            batch, rel_length52        )53        return predicted_words[0]54 55    def encode_batch(self, wavs):56        """Encodes the input audio into a sequence of hidden states57        The waveforms should already be in the model's desired format.58        You can call:59        ``normalized = EncoderDecoderASR.normalizer(signal, sample_rate)``60        to get a correctly converted signal in most cases.61        Arguments62        ---------63        wavs : torch.tensor64            Batch of waveforms [batch, time, channels] or [batch, time]65            depending on the model.66        wav_lens : torch.tensor67            Lengths of the waveforms relative to the longest one in the68            batch, tensor of shape [batch]. The longest one should have69            relative length 1.0 and others len(waveform) / max_length.70            Used for ignoring padding.71        Returns72        -------73        torch.tensor74            The encoded batch75        """76        wavs = wavs.float()77        wavs = wavs.to(self.device)78        outputs = self.mods.wav2vec2(wavs)79        outputs = self.mods.enc(outputs)80        outputs = self.mods.ctc_lin(outputs)81        return outputs82 83    def transcribe_batch(self, wavs, wav_lens):84        """Transcribes the input audio into a sequence of words85        The waveforms should already be in the model's desired format.86        You can call:87        ``normalized = EncoderDecoderASR.normalizer(signal, sample_rate)``88        to get a correctly converted signal in most cases.89        Arguments90        ---------91        wavs : torch.tensor92            Batch of waveforms [batch, time, channels] or [batch, time]93            depending on the model.94        wav_lens : torch.tensor95            Lengths of the waveforms relative to the longest one in the96            batch, tensor of shape [batch]. The longest one should have97            relative length 1.0 and others len(waveform) / max_length.98            Used for ignoring padding.99        Returns100        -------101        list102            Each waveform in the batch transcribed.103        tensor104            Each predicted token id.105        """106        with torch.no_grad():107            wav_lens = wav_lens.to(self.device)108            encoder_out = self.encode_batch(wavs)109            p_ctc = self.hparams.log_softmax(encoder_out)110            sequences = self.hparams.decoder(p_ctc, wav_lens)111            predicted_words_list = []112            for sequence in sequences:113                predicted_tokens = self.tokenizer.convert_ids_to_tokens(114                    sequence115                )116                predicted_words = []117                for c in predicted_tokens:118                    if c == "[CLS]":119                        continue120                    elif c == "[SEP]" or c == "[PAD]":121                        break122                    else:123                        predicted_words.append(c)124                predicted_words_list.append(predicted_words)125 126        return predicted_words_list127 128    def forward(self, wavs, wav_lens):129        """Runs full transcription - note: no gradients through decoding"""130        return self.transcribe_batch(wavs, wav_lens)131