torchaudio 2.8.0__cp310-cp310-win_amd64.whl → 2.9.0__cp310-cp310-win_amd64.whl

torchaudio/prototype/models/rnnt_decoder.py

@@ -1,402 +0,0 @@
-from typing import Callable, Dict, List, Optional, Tuple
-
-import torch
-from torchaudio.models import RNNT
-from torchaudio.prototype.models.rnnt import TrieNode
-
-from torchaudio._internal.module_utils import dropping_class_support
-
-__all__ = ["Hypothesis", "RNNTBeamSearchBiasing"]
-
-
-Hypothesis = Tuple[List[int], torch.Tensor, List[List[torch.Tensor]], float, list]
-Hypothesis.__doc__ = """Hypothesis generated by RNN-T beam search decoder,
-    represented as tuple of (tokens, prediction network output, prediction network state, score).
-    """
-
-
-def _get_hypo_tokens(hypo: Hypothesis) -> List[int]:
-    return hypo[0]
-
-
-def _get_hypo_predictor_out(hypo: Hypothesis) -> torch.Tensor:
-    return hypo[1]
-
-
-def _get_hypo_state(hypo: Hypothesis) -> List[List[torch.Tensor]]:
-    return hypo[2]
-
-
-def _get_hypo_score(hypo: Hypothesis) -> float:
-    return hypo[3]
-
-
-def _get_hypo_trie(hypo: Hypothesis) -> TrieNode:
-    return hypo[4]
-
-
-def _set_hypo_trie(hypo: Hypothesis, trie: TrieNode) -> None:
-    hypo[4] = trie
-
-
-def _get_hypo_key(hypo: Hypothesis) -> str:
-    return str(hypo[0])
-
-
-def _batch_state(hypos: List[Hypothesis]) -> List[List[torch.Tensor]]:
-    states: List[List[torch.Tensor]] = []
-    for i in range(len(_get_hypo_state(hypos[0]))):
-        batched_state_components: List[torch.Tensor] = []
-        for j in range(len(_get_hypo_state(hypos[0])[i])):
-            batched_state_components.append(torch.cat([_get_hypo_state(hypo)[i][j] for hypo in hypos]))
-        states.append(batched_state_components)
-    return states
-
-
-def _slice_state(states: List[List[torch.Tensor]], idx: int, device: torch.device) -> List[List[torch.Tensor]]:
-    idx_tensor = torch.tensor([idx], device=device)
-    return [[state.index_select(0, idx_tensor) for state in state_tuple] for state_tuple in states]
-
-
-def _default_hypo_sort_key(hypo: Hypothesis) -> float:
-    return _get_hypo_score(hypo) / (len(_get_hypo_tokens(hypo)) + 1)
-
-
-def _compute_updated_scores(
-    hypos: List[Hypothesis],
-    next_token_probs: torch.Tensor,
-    beam_width: int,
-) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    hypo_scores = torch.tensor([_get_hypo_score(h) for h in hypos]).unsqueeze(1)
-    nonblank_scores = hypo_scores + next_token_probs[:, :-1]  # [beam_width, num_tokens - 1]
-    nonblank_nbest_scores, nonblank_nbest_idx = nonblank_scores.reshape(-1).topk(beam_width)
-    nonblank_nbest_hypo_idx = nonblank_nbest_idx.div(nonblank_scores.shape[1], rounding_mode="trunc")
-    nonblank_nbest_token = nonblank_nbest_idx % nonblank_scores.shape[1]
-    return nonblank_nbest_scores, nonblank_nbest_hypo_idx, nonblank_nbest_token
-
-
-def _remove_hypo(hypo: Hypothesis, hypo_list: List[Hypothesis]) -> None:
-    for i, elem in enumerate(hypo_list):
-        if _get_hypo_key(hypo) == _get_hypo_key(elem):
-            del hypo_list[i]
-            break
-
-
-@dropping_class_support
-class RNNTBeamSearchBiasing(torch.nn.Module):
-    r"""Beam search decoder for RNN-T model with biasing support.
-
-    Args:
-        model (RNNT): RNN-T model to use.
-        blank (int): index of blank token in vocabulary.
-        temperature (float, optional): temperature to apply to joint network output.
-            Larger values yield more uniform samples. (Default: 1.0)
-        hypo_sort_key (Callable[[Hypothesis], float] or None, optional): callable that computes a score
-            for a given hypothesis to rank hypotheses by. If ``None``, defaults to callable that returns
-            hypothesis score normalized by token sequence length. (Default: None)
-        step_max_tokens (int, optional): maximum number of tokens to emit per input time step. (Default: 100)
-        trie (list, optional): the prefix tree for TCPGen biasing
-        biasing (bool, optional): If true, do biasing, otherwise use standard RNN-T support
-    """
-
-    def __init__(
-        self,
-        model: RNNT,
-        blank: int,
-        temperature: float = 1.0,
-        hypo_sort_key: Optional[Callable[[Hypothesis], float]] = None,
-        step_max_tokens: int = 100,
-        trie: TrieNode = None,
-        biasing: bool = False,
-    ) -> None:
-        super().__init__()
-        self.model = model
-        self.blank = blank
-        self.temperature = temperature
-        self.resettrie = trie or []
-        self.dobiasing = biasing
-
-        if hypo_sort_key is None:
-            self.hypo_sort_key = _default_hypo_sort_key
-        else:
-            self.hypo_sort_key = hypo_sort_key
-
-        self.step_max_tokens = step_max_tokens
-
-    def _init_b_hypos(self, hypo: Optional[Hypothesis], device: torch.device) -> List[Hypothesis]:
-        if hypo is not None:
-            token = _get_hypo_tokens(hypo)[-1]
-            state = _get_hypo_state(hypo)
-        else:
-            token = self.blank
-            state = None
-
-        one_tensor = torch.tensor([1], device=device)
-        pred_out, _, pred_state = self.model.predict(torch.tensor([[token]], device=device), one_tensor, state)
-        init_hypo = ([token], pred_out[0].detach(), pred_state, 0.0, self.resettrie)
-        return [init_hypo]
-
-    def _get_trie_mask(self, trie):
-        step_mask = torch.ones(len(self.model.char_list) + 1)
-        step_mask[list(trie[0].keys())] = 0
-        # step_mask[-1] = 0
-        return step_mask
-
-    def _get_generation_prob(self, trie):
-        if len(trie[0].keys()) == 0:
-            return True
-        else:
-            return False
-
-    def _gen_next_token_probs(
-        self, enc_out: torch.Tensor, hypos: List[Hypothesis], device: torch.device
-    ) -> torch.Tensor:
-        one_tensor = torch.tensor([1], device=device)
-        predictor_out = torch.stack([_get_hypo_predictor_out(h) for h in hypos], dim=0)
-        if self.dobiasing:
-            # Get valid subset of wordpieces
-            trie_masks = torch.stack([self._get_trie_mask(_get_hypo_trie(h)) for h in hypos], dim=0)
-            trie_masks = trie_masks.to(enc_out.device).unsqueeze(1)  # beam_width, 1, nchars
-            # Determine if there is any paths on the trie
-            genprob_masks = torch.tensor([self._get_generation_prob(_get_hypo_trie(h)) for h in hypos])  # beam_width
-            genprob_masks = genprob_masks.to(enc_out.device)
-            # Forward TCPGen component
-            last_tokens = torch.tensor([_get_hypo_tokens(h)[-1] for h in hypos]).unsqueeze(-1).to(enc_out.device)
-            hptr, tcpgen_dist = self.model.forward_tcpgen(last_tokens, trie_masks, enc_out)
-        else:
-            hptr = None
-        # hptr sent to joiner, if deepbiasing is True joiner will use it
-        joined_out, _, joined_activation = self.model.join(
-            enc_out,
-            one_tensor,
-            predictor_out,
-            torch.tensor([1] * len(hypos), device=device),
-            hptr=hptr,
-        )  # [beam_width, 1, 1, num_tokens]
-        if self.dobiasing:
-            p_gen = torch.sigmoid(self.model.pointer_gate(torch.cat((joined_activation, hptr), dim=-1)))
-            p_gen = p_gen.masked_fill(genprob_masks.view(p_gen.size(0), 1, 1, 1), 0)
-            model_tu = torch.softmax(joined_out / self.temperature, dim=3)
-            # assuming last token is blank
-            p_not_null = 1.0 - model_tu[:, :, :, -1:]
-            ptr_dist_fact = torch.cat([tcpgen_dist[:, :, :, :-2], tcpgen_dist[:, :, :, -1:]], dim=-1) * p_not_null
-            ptr_gen_complement = tcpgen_dist[:, :, :, -1:] * p_gen
-            p_partial = ptr_dist_fact[:, :, :, :-1] * p_gen + model_tu[:, :, :, :-1] * (1 - p_gen + ptr_gen_complement)
-            p_final = torch.cat([p_partial, model_tu[:, :, :, -1:]], dim=-1)
-            joined_out = torch.log(p_final)
-        else:
-            joined_out = torch.nn.functional.log_softmax(joined_out / self.temperature, dim=3)
-        return joined_out[:, 0, 0]
-
-    def _gen_b_hypos(
-        self,
-        b_hypos: List[Hypothesis],
-        a_hypos: List[Hypothesis],
-        next_token_probs: torch.Tensor,
-        key_to_b_hypo: Dict[str, Hypothesis],
-    ) -> List[Hypothesis]:
-        for i in range(len(a_hypos)):
-            h_a = a_hypos[i]
-            append_blank_score = _get_hypo_score(h_a) + next_token_probs[i, -1]
-            if _get_hypo_key(h_a) in key_to_b_hypo:
-                h_b = key_to_b_hypo[_get_hypo_key(h_a)]
-                _remove_hypo(h_b, b_hypos)
-                score = float(torch.tensor(_get_hypo_score(h_b)).logaddexp(append_blank_score))
-            else:
-                score = float(append_blank_score)
-            h_b = (
-                _get_hypo_tokens(h_a),
-                _get_hypo_predictor_out(h_a),
-                _get_hypo_state(h_a),
-                score,
-                _get_hypo_trie(h_a),
-            )
-            b_hypos.append(h_b)
-            key_to_b_hypo[_get_hypo_key(h_b)] = h_b
-        _, sorted_idx = torch.tensor([_get_hypo_score(hypo) for hypo in b_hypos]).sort()
-        return [b_hypos[idx] for idx in sorted_idx]
-
-    def _gen_a_hypos(
-        self,
-        a_hypos: List[Hypothesis],
-        b_hypos: List[Hypothesis],
-        next_token_probs: torch.Tensor,
-        t: int,
-        beam_width: int,
-        device: torch.device,
-    ) -> List[Hypothesis]:
-        (
-            nonblank_nbest_scores,
-            nonblank_nbest_hypo_idx,
-            nonblank_nbest_token,
-        ) = _compute_updated_scores(a_hypos, next_token_probs, beam_width)
-
-        if len(b_hypos) < beam_width:
-            b_nbest_score = -float("inf")
-        else:
-            b_nbest_score = _get_hypo_score(b_hypos[-beam_width])
-
-        base_hypos: List[Hypothesis] = []
-        new_tokens: List[int] = []
-        new_scores: List[float] = []
-        for i in range(beam_width):
-            score = float(nonblank_nbest_scores[i])
-            if score > b_nbest_score:
-                a_hypo_idx = int(nonblank_nbest_hypo_idx[i])
-                base_hypos.append(a_hypos[a_hypo_idx])
-                new_tokens.append(int(nonblank_nbest_token[i]))
-                new_scores.append(score)
-
-        if base_hypos:
-            new_hypos = self._gen_new_hypos(base_hypos, new_tokens, new_scores, t, device)
-        else:
-            new_hypos: List[Hypothesis] = []
-
-        return new_hypos
-
-    def _gen_new_hypos(
-        self,
-        base_hypos: List[Hypothesis],
-        tokens: List[int],
-        scores: List[float],
-        t: int,
-        device: torch.device,
-    ) -> List[Hypothesis]:
-        tgt_tokens = torch.tensor([[token] for token in tokens], device=device)
-        states = _batch_state(base_hypos)
-        pred_out, _, pred_states = self.model.predict(
-            tgt_tokens,
-            torch.tensor([1] * len(base_hypos), device=device),
-            states,
-        )
-        new_hypos: List[Hypothesis] = []
-        for i, h_a in enumerate(base_hypos):
-            new_tokens = _get_hypo_tokens(h_a) + [tokens[i]]
-            if self.dobiasing:
-                new_trie = self.model.get_tcpgen_step(tokens[i], _get_hypo_trie(h_a), self.resettrie)
-            else:
-                new_trie = self.resettrie
-            new_hypos.append(
-                (new_tokens, pred_out[i].detach(), _slice_state(pred_states, i, device), scores[i], new_trie)
-            )
-        return new_hypos
-
-    def _search(
-        self,
-        enc_out: torch.Tensor,
-        hypo: Optional[Hypothesis],
-        beam_width: int,
-    ) -> List[Hypothesis]:
-        n_time_steps = enc_out.shape[1]
-        device = enc_out.device
-
-        a_hypos: List[Hypothesis] = []
-        b_hypos = self._init_b_hypos(hypo, device)
-        for t in range(n_time_steps):
-            a_hypos = b_hypos
-            b_hypos = torch.jit.annotate(List[Hypothesis], [])
-            key_to_b_hypo: Dict[str, Hypothesis] = {}
-            symbols_current_t = 0
-
-            while a_hypos:
-                next_token_probs = self._gen_next_token_probs(enc_out[:, t : t + 1], a_hypos, device)
-                next_token_probs = next_token_probs.cpu()
-                b_hypos = self._gen_b_hypos(b_hypos, a_hypos, next_token_probs, key_to_b_hypo)
-
-                if symbols_current_t == self.step_max_tokens:
-                    break
-
-                a_hypos = self._gen_a_hypos(
-                    a_hypos,
-                    b_hypos,
-                    next_token_probs,
-                    t,
-                    beam_width,
-                    device,
-                )
-                if a_hypos:
-                    symbols_current_t += 1
-
-            _, sorted_idx = torch.tensor([self.hypo_sort_key(hypo) for hypo in b_hypos]).topk(beam_width)
-            b_hypos = [b_hypos[idx] for idx in sorted_idx]
-
-        return b_hypos
-
-    def forward(
-        self,
-        input: torch.Tensor,
-        length: torch.Tensor,
-        beam_width: int,
-    ) -> List[Hypothesis]:
-        r"""Performs beam search for the given input sequence.
-
-        T: number of frames;
-        D: feature dimension of each frame.
-
-        Args:
-            input (torch.Tensor): sequence of input frames, with shape (T, D) or (1, T, D).
-            length (torch.Tensor): number of valid frames in input
-                sequence, with shape () or (1,).
-            beam_width (int): beam size to use during search.
-
-        Returns:
-            List[Hypothesis]: top-``beam_width`` hypotheses found by beam search.
-        """
-        if input.dim() != 2 and not (input.dim() == 3 and input.shape[0] == 1):
-            raise ValueError("input must be of shape (T, D) or (1, T, D)")
-        if input.dim() == 2:
-            input = input.unsqueeze(0)
-
-        if length.shape != () and length.shape != (1,):
-            raise ValueError("length must be of shape () or (1,)")
-        if input.dim() == 0:
-            input = input.unsqueeze(0)
-
-        enc_out, _ = self.model.transcribe(input, length)
-        return self._search(enc_out, None, beam_width)
-
-    @torch.jit.export
-    def infer(
-        self,
-        input: torch.Tensor,
-        length: torch.Tensor,
-        beam_width: int,
-        state: Optional[List[List[torch.Tensor]]] = None,
-        hypothesis: Optional[Hypothesis] = None,
-    ) -> Tuple[List[Hypothesis], List[List[torch.Tensor]]]:
-        r"""Performs beam search for the given input sequence in streaming mode.
-
-        T: number of frames;
-        D: feature dimension of each frame.
-
-        Args:
-            input (torch.Tensor): sequence of input frames, with shape (T, D) or (1, T, D).
-            length (torch.Tensor): number of valid frames in input
-                sequence, with shape () or (1,).
-            beam_width (int): beam size to use during search.
-            state (List[List[torch.Tensor]] or None, optional): list of lists of tensors
-                representing transcription network internal state generated in preceding
-                invocation. (Default: ``None``)
-            hypothesis (Hypothesis or None): hypothesis from preceding invocation to seed
-                search with. (Default: ``None``)
-
-        Returns:
-            (List[Hypothesis], List[List[torch.Tensor]]):
-                List[Hypothesis]
-                    top-``beam_width`` hypotheses found by beam search.
-                List[List[torch.Tensor]]
-                    list of lists of tensors representing transcription network
-                    internal state generated in current invocation.
-        """
-        if input.dim() != 2 and not (input.dim() == 3 and input.shape[0] == 1):
-            raise ValueError("input must be of shape (T, D) or (1, T, D)")
-        if input.dim() == 2:
-            input = input.unsqueeze(0)
-
-        if length.shape != () and length.shape != (1,):
-            raise ValueError("length must be of shape () or (1,)")
-        if length.dim() == 0:
-            length = length.unsqueeze(0)
-
-        enc_out, _, state = self.model.transcribe_streaming(input, length, state)
-        return self._search(enc_out, hypothesis, beam_width), state

torchaudio/prototype/pipelines/__init__.py

@@ -1,21 +0,0 @@
-from ._vggish import VGGISH, VGGishBundle
-from .hifigan_pipeline import HIFIGAN_VOCODER_V3_LJSPEECH as _HIFIGAN_VOCODER_V3_LJSPEECH, HiFiGANVocoderBundle
-from .rnnt_pipeline import (
-    EMFORMER_RNNT_BASE_MUSTC as _EMFORMER_RNNT_BASE_MUSTC,
-    EMFORMER_RNNT_BASE_TEDLIUM3 as _EMFORMER_RNNT_BASE_TEDLIUM3
-)
-from torchaudio._internal.module_utils import dropping_const_support
-
-EMFORMER_RNNT_BASE_MUSTC = dropping_const_support(_EMFORMER_RNNT_BASE_MUSTC)
-EMFORMER_RNNT_BASE_TEDLIUM3 = dropping_const_support(_EMFORMER_RNNT_BASE_TEDLIUM3)
-HIFIGAN_VOCODER_V3_LJSPEECH = dropping_const_support(_HIFIGAN_VOCODER_V3_LJSPEECH)
-
-
-__all__ = [
-    "EMFORMER_RNNT_BASE_MUSTC",
-    "EMFORMER_RNNT_BASE_TEDLIUM3",
-    "HIFIGAN_VOCODER_V3_LJSPEECH",
-    "HiFiGANVocoderBundle",
-    "VGGISH",
-    "VGGishBundle",
-]

torchaudio/prototype/pipelines/_vggish/__init__.py

@@ -1,7 +0,0 @@
-from ._vggish_pipeline import VGGISH as _VGGISH, VGGishBundle
-from torchaudio._internal.module_utils import dropping_const_support
-
-
-VGGISH = dropping_const_support(_VGGISH, "VGGISH")
-
-__all__ = ["VGGISH", "VGGishBundle"]

torchaudio/prototype/pipelines/_vggish/_vggish_impl.py

@@ -1,236 +0,0 @@
-# Derived from torchvggish (https://github.com/harritaylor/torchvggish).
-# Copyright 2017 The TensorFlow Authors All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-import math
-
-import torch
-
-from torchaudio._internal.module_utils import dropping_class_support
-
-
-_MEL_BREAK_FREQUENCY_HERTZ = 700.0
-_MEL_HIGH_FREQUENCY_Q = 1127.0
-
-
-_SAMPLE_RATE = 16000
-_STFT_WINDOW_LENGTH_SECONDS = 0.025
-_STFT_HOP_LENGTH_SECONDS = 0.010
-_MEL_MIN_HZ = 125
-_MEL_MAX_HZ = 7500
-_NUM_BANDS = 64
-_LOG_OFFSET = 0.01
-_EXAMPLE_WINDOW_SECONDS = 0.96  # Each example contains 96 10ms frames
-_EXAMPLE_HOP_SECONDS = 0.96  # with zero overlap.
-
-
-def _build_features_network():
-    layers = []
-
-    for input_dim, output_dim in [(1, 64), (64, 128)]:
-        layers += [
-            torch.nn.Conv2d(input_dim, output_dim, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
-            torch.nn.ReLU(inplace=True),
-            torch.nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False),
-        ]
-
-    for input_dim, output_dim in [(128, 256), (256, 512)]:
-        layers += [
-            torch.nn.Conv2d(input_dim, output_dim, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
-            torch.nn.ReLU(inplace=True),
-            torch.nn.Conv2d(
-                output_dim,
-                output_dim,
-                kernel_size=(3, 3),
-                stride=(1, 1),
-                padding=(1, 1),
-            ),
-            torch.nn.ReLU(inplace=True),
-            torch.nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False),
-        ]
-
-    return torch.nn.Sequential(*layers)
-
-
-def _build_embedding_network():
-    return torch.nn.Sequential(
-        torch.nn.Linear(512 * 4 * 6, 4096),
-        torch.nn.ReLU(True),
-        torch.nn.Linear(4096, 4096),
-        torch.nn.ReLU(True),
-        torch.nn.Linear(4096, 128),
-        torch.nn.ReLU(True),
-    )
-
-
-def _frame(data, window_length, hop_length):
-    num_samples = data.shape[0]
-    num_frames = 1 + int(math.floor((num_samples - window_length) / hop_length))
-    shape = (num_frames, window_length) + data.shape[1:]
-    strides = (data.stride()[0] * hop_length,) + data.stride()
-    return torch.as_strided(data, shape, strides)
-
-
-def _stft_magnitude(signal, fft_length, hop_length=None, window_length=None):
-    frames = _frame(signal, window_length, hop_length)
-    window = torch.hann_window(window_length, periodic=True).to(signal.device)
-    windowed_frames = frames * window
-    return torch.abs(torch.fft.rfft(windowed_frames, int(fft_length)))
-
-
-def _hertz_to_mel(frequencies_hertz):
-    return _MEL_HIGH_FREQUENCY_Q * torch.log(1.0 + (frequencies_hertz / _MEL_BREAK_FREQUENCY_HERTZ))
-
-
-def _spectrogram_to_mel_matrix(
-    num_mel_bins=20,
-    num_spectrogram_bins=129,
-    audio_sample_rate=8000,
-    lower_edge_hertz=125.0,
-    upper_edge_hertz=3800.0,
-):
-    nyquist_hertz = audio_sample_rate / 2.0
-    if lower_edge_hertz < 0.0:
-        raise ValueError("lower_edge_hertz %.1f must be >= 0" % lower_edge_hertz)
-    if lower_edge_hertz >= upper_edge_hertz:
-        raise ValueError("lower_edge_hertz %.1f >= upper_edge_hertz %.1f" % (lower_edge_hertz, upper_edge_hertz))
-
-    if upper_edge_hertz > nyquist_hertz:
-        raise ValueError("upper_edge_hertz %.1f is greater than Nyquist %.1f" % (upper_edge_hertz, nyquist_hertz))
-    spectrogram_bins_hertz = torch.linspace(0.0, nyquist_hertz, num_spectrogram_bins)
-
-    spectrogram_bins_mel = _hertz_to_mel(spectrogram_bins_hertz)
-    # The i'th mel band (starting from i=1) has center frequency
-    # band_edges_mel[i], lower edge band_edges_mel[i-1], and higher edge
-    # band_edges_mel[i+1].  Thus, we need num_mel_bins + 2 values in
-    # the band_edges_mel arrays.
-    band_edges_mel = torch.linspace(
-        _hertz_to_mel(torch.tensor(lower_edge_hertz)),
-        _hertz_to_mel(torch.tensor(upper_edge_hertz)),
-        num_mel_bins + 2,
-    )
-    # Matrix to post-multiply feature arrays whose rows are num_spectrogram_bins
-    # of spectrogram values.
-    mel_weights_matrix = torch.empty((num_spectrogram_bins, num_mel_bins))
-    for i in range(num_mel_bins):
-        lower_edge_mel, center_mel, upper_edge_mel = band_edges_mel[i : i + 3]
-        # Calculate lower and upper slopes for every spectrogram bin.
-        # Line segments are linear in the *mel* domain, not hertz.
-        lower_slope = (spectrogram_bins_mel - lower_edge_mel) / (center_mel - lower_edge_mel)
-        upper_slope = (upper_edge_mel - spectrogram_bins_mel) / (upper_edge_mel - center_mel)
-
-        # .. then intersect them with each other and zero.
-        mel_weights_matrix[:, i] = torch.maximum(torch.tensor(0.0), torch.minimum(lower_slope, upper_slope))
-
-    # HTK excludes the spectrogram DC bin; make sure it always gets a zero
-    # coefficient.
-    mel_weights_matrix[0, :] = 0.0
-    return mel_weights_matrix
-
-
-def _log_mel_spectrogram(
-    data,
-    audio_sample_rate=8000,
-    log_offset=0.0,
-    window_length_secs=0.025,
-    hop_length_secs=0.010,
-    **kwargs,
-):
-    window_length_samples = int(round(audio_sample_rate * window_length_secs))
-    hop_length_samples = int(round(audio_sample_rate * hop_length_secs))
-    fft_length = 2 ** int(math.ceil(math.log(window_length_samples) / math.log(2.0)))
-
-    spectrogram = _stft_magnitude(
-        data,
-        fft_length=fft_length,
-        hop_length=hop_length_samples,
-        window_length=window_length_samples,
-    )
-    mel_spectrogram = torch.matmul(
-        spectrogram,
-        _spectrogram_to_mel_matrix(
-            num_spectrogram_bins=spectrogram.shape[1],
-            audio_sample_rate=audio_sample_rate,
-            **kwargs,
-        ).to(spectrogram),
-    )
-    return torch.log(mel_spectrogram + log_offset)
-
-
-def _waveform_to_examples(data):
-    # Compute log mel spectrogram features, with shape (n_frame, n_mel)
-    log_mel = _log_mel_spectrogram(
-        data,
-        audio_sample_rate=_SAMPLE_RATE,
-        log_offset=_LOG_OFFSET,
-        window_length_secs=_STFT_WINDOW_LENGTH_SECONDS,
-        hop_length_secs=_STFT_HOP_LENGTH_SECONDS,
-        num_mel_bins=_NUM_BANDS,
-        lower_edge_hertz=_MEL_MIN_HZ,
-        upper_edge_hertz=_MEL_MAX_HZ,
-    )
-
-    # Frame features into examples, with shape (n_example, n_frame, n_mel)
-    features_sample_rate = 1.0 / _STFT_HOP_LENGTH_SECONDS
-    example_window_length = int(round(_EXAMPLE_WINDOW_SECONDS * features_sample_rate))
-
-    example_hop_length = int(round(_EXAMPLE_HOP_SECONDS * features_sample_rate))
-    log_mel_examples = _frame(log_mel, window_length=example_window_length, hop_length=example_hop_length)
-
-    # (n_example, 1, n_frame, n_mel)
-    return log_mel_examples.unsqueeze(1)
-
-
-@dropping_class_support
-class VGGish(torch.nn.Module):
-    """Implementation of VGGish model :cite:`45611`."""
-
-    def __init__(self):
-        super().__init__()
-
-        self.features_network = _build_features_network()
-        self.embedding_network = _build_embedding_network()
-
-    def forward(self, input: torch.Tensor) -> torch.Tensor:
-        """
-        Args:
-            input (torch.Tensor): batch of spectrograms, with shape `(n_example, 1, n_frame, 64)`.
-
-        Returns:
-            torch.Tensor: model output, with shape `(n_example, 128)`.
-        """
-        x = self.features_network(input)
-
-        x = x.permute(0, 2, 3, 1)
-        x = x.reshape(x.size(0), -1)
-
-        return self.embedding_network(x)
-
-@dropping_class_support
-class VGGishInputProcessor:
-    """Converts raw waveforms to batches of examples to use as inputs to VGGish."""
-
-    def __call__(self, input: torch.Tensor) -> torch.Tensor:
-        """
-        Args:
-            input (torch.Tensor): waveform, with shape `(T,)`.
-            sample_rate (int): sample rate of waveform in hertz.
-
-        Returns:
-            torch.Tensor: batch of examples to pass to VGGish, with shape `(n_example, 1, n_frame, 64)`.
-        """
-        if len(input.shape) != 1:
-            raise ValueError("input waveform must have dimension of 1.")
-        return _waveform_to_examples(input)