torchaudio 2.9.0__cp314-cp314-macosx_11_0_arm64.whl

torchaudio/pipelines/rnnt_pipeline.py

@@ -0,0 +1,380 @@
+import json
+import math
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from functools import partial
+from typing import Callable, List, Tuple
+
+import torch
+import torchaudio
+from torchaudio._internal import module_utils
+from torchaudio.models import emformer_rnnt_base, RNNT, RNNTBeamSearch
+
+
+__all__ = []
+
+_decibel = 2 * 20 * math.log10(torch.iinfo(torch.int16).max)
+_gain = pow(10, 0.05 * _decibel)
+
+
+def _piecewise_linear_log(x):
+    x[x > math.e] = torch.log(x[x > math.e])
+    x[x <= math.e] = x[x <= math.e] / math.e
+    return x
+
+
+class _FunctionalModule(torch.nn.Module):
+    def __init__(self, functional):
+        super().__init__()
+        self.functional = functional
+
+    def forward(self, input):
+        return self.functional(input)
+
+
+class _GlobalStatsNormalization(torch.nn.Module):
+    def __init__(self, global_stats_path):
+        super().__init__()
+
+        with open(global_stats_path) as f:
+            blob = json.loads(f.read())
+
+        self.register_buffer("mean", torch.tensor(blob["mean"]))
+        self.register_buffer("invstddev", torch.tensor(blob["invstddev"]))
+
+    def forward(self, input):
+        return (input - self.mean) * self.invstddev
+
+
+class _FeatureExtractor(ABC):
+    @abstractmethod
+    def __call__(self, input: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Generates features and length output from the given input tensor.
+
+        Args:
+            input (torch.Tensor): input tensor.
+
+        Returns:
+            (torch.Tensor, torch.Tensor):
+            torch.Tensor:
+                Features, with shape `(length, *)`.
+            torch.Tensor:
+                Length, with shape `(1,)`.
+        """
+
+
+class _TokenProcessor(ABC):
+    @abstractmethod
+    def __call__(self, tokens: List[int], **kwargs) -> str:
+        """Decodes given list of tokens to text sequence.
+
+        Args:
+            tokens (List[int]): list of tokens to decode.
+
+        Returns:
+            str:
+                Decoded text sequence.
+        """
+
+
+class _ModuleFeatureExtractor(torch.nn.Module, _FeatureExtractor):
+    """``torch.nn.Module``-based feature extraction pipeline.
+
+    Args:
+        pipeline (torch.nn.Module): module that implements feature extraction logic.
+    """
+
+    def __init__(self, pipeline: torch.nn.Module) -> None:
+        super().__init__()
+        self.pipeline = pipeline
+
+    def forward(self, input: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Generates features and length output from the given input tensor.
+
+        Args:
+            input (torch.Tensor): input tensor.
+
+        Returns:
+            (torch.Tensor, torch.Tensor):
+            torch.Tensor:
+                Features, with shape `(length, *)`.
+            torch.Tensor:
+                Length, with shape `(1,)`.
+        """
+        features = self.pipeline(input)
+        length = torch.tensor([features.shape[0]])
+        return features, length
+
+
+class _SentencePieceTokenProcessor(_TokenProcessor):
+    """SentencePiece-model-based token processor.
+
+    Args:
+        sp_model_path (str): path to SentencePiece model.
+    """
+
+    def __init__(self, sp_model_path: str) -> None:
+        if not module_utils.is_module_available("sentencepiece"):
+            raise RuntimeError("SentencePiece is not available. Please install it.")
+
+        import sentencepiece as spm
+
+        self.sp_model = spm.SentencePieceProcessor(model_file=sp_model_path)
+        self.post_process_remove_list = {
+            self.sp_model.unk_id(),
+            self.sp_model.eos_id(),
+            self.sp_model.pad_id(),
+        }
+
+    def __call__(self, tokens: List[int], lstrip: bool = True) -> str:
+        """Decodes given list of tokens to text sequence.
+
+        Args:
+            tokens (List[int]): list of tokens to decode.
+            lstrip (bool, optional): if ``True``, returns text sequence with leading whitespace
+                removed. (Default: ``True``).
+
+        Returns:
+            str:
+                Decoded text sequence.
+        """
+        filtered_hypo_tokens = [
+            token_index for token_index in tokens[1:] if token_index not in self.post_process_remove_list
+        ]
+        output_string = "".join(self.sp_model.id_to_piece(filtered_hypo_tokens)).replace("\u2581", " ")
+
+        if lstrip:
+            return output_string.lstrip()
+        else:
+            return output_string
+
+
+@dataclass
+class RNNTBundle:
+    """Dataclass that bundles components for performing automatic speech recognition (ASR, speech-to-text)
+    inference with an RNN-T model.
+
+    More specifically, the class provides methods that produce the featurization pipeline,
+    decoder wrapping the specified RNN-T model, and output token post-processor that together
+    constitute a complete end-to-end ASR inference pipeline that produces a text sequence
+    given a raw waveform.
+
+    It can support non-streaming (full-context) inference as well as streaming inference.
+
+    Users should not directly instantiate objects of this class; rather, users should use the
+    instances (representing pre-trained models) that exist within the module,
+    e.g. :data:`torchaudio.pipelines.EMFORMER_RNNT_BASE_LIBRISPEECH`.
+
+    Example
+        >>> import torchaudio
+        >>> from torchaudio.pipelines import EMFORMER_RNNT_BASE_LIBRISPEECH
+        >>> import torch
+        >>>
+        >>> # Non-streaming inference.
+        >>> # Build feature extractor, decoder with RNN-T model, and token processor.
+        >>> feature_extractor = EMFORMER_RNNT_BASE_LIBRISPEECH.get_feature_extractor()
+        100%|███████████████████████████████| 3.81k/3.81k [00:00<00:00, 4.22MB/s]
+        >>> decoder = EMFORMER_RNNT_BASE_LIBRISPEECH.get_decoder()
+        Downloading: "https://download.pytorch.org/torchaudio/models/emformer_rnnt_base_librispeech.pt"
+        100%|███████████████████████████████| 293M/293M [00:07<00:00, 42.1MB/s]
+        >>> token_processor = EMFORMER_RNNT_BASE_LIBRISPEECH.get_token_processor()
+        100%|███████████████████████████████| 295k/295k [00:00<00:00, 25.4MB/s]
+        >>>
+        >>> # Instantiate LibriSpeech dataset; retrieve waveform for first sample.
+        >>> dataset = torchaudio.datasets.LIBRISPEECH("/home/librispeech", url="test-clean")
+        >>> waveform = next(iter(dataset))[0].squeeze()
+        >>>
+        >>> with torch.no_grad():
+        >>>     # Produce mel-scale spectrogram features.
+        >>>     features, length = feature_extractor(waveform)
+        >>>
+        >>>     # Generate top-10 hypotheses.
+        >>>     hypotheses = decoder(features, length, 10)
+        >>>
+        >>> # For top hypothesis, convert predicted tokens to text.
+        >>> text = token_processor(hypotheses[0][0])
+        >>> print(text)
+        he hoped there would be stew for dinner turnips and carrots and bruised potatoes and fat mutton pieces to [...]
+        >>>
+        >>>
+        >>> # Streaming inference.
+        >>> hop_length = EMFORMER_RNNT_BASE_LIBRISPEECH.hop_length
+        >>> num_samples_segment = EMFORMER_RNNT_BASE_LIBRISPEECH.segment_length * hop_length
+        >>> num_samples_segment_right_context = (
+        >>>     num_samples_segment + EMFORMER_RNNT_BASE_LIBRISPEECH.right_context_length * hop_length
+        >>> )
+        >>>
+        >>> # Build streaming inference feature extractor.
+        >>> streaming_feature_extractor = EMFORMER_RNNT_BASE_LIBRISPEECH.get_streaming_feature_extractor()
+        >>>
+        >>> # Process same waveform as before, this time sequentially across overlapping segments
+        >>> # to simulate streaming inference. Note the usage of ``streaming_feature_extractor`` and ``decoder.infer``.
+        >>> state, hypothesis = None, None
+        >>> for idx in range(0, len(waveform), num_samples_segment):
+        >>>     segment = waveform[idx: idx + num_samples_segment_right_context]
+        >>>     segment = torch.nn.functional.pad(segment, (0, num_samples_segment_right_context - len(segment)))
+        >>>     with torch.no_grad():
+        >>>         features, length = streaming_feature_extractor(segment)
+        >>>         hypotheses, state = decoder.infer(features, length, 10, state=state, hypothesis=hypothesis)
+        >>>     hypothesis = hypotheses[0]
+        >>>     transcript = token_processor(hypothesis[0])
+        >>>     if transcript:
+        >>>         print(transcript, end=" ", flush=True)
+        he hoped there would be stew for dinner turn ips and car rots and bru 'd oes and fat mut ton pieces to [...]
+    """
+
+    class FeatureExtractor(_FeatureExtractor):
+        """Interface of the feature extraction part of RNN-T pipeline"""
+
+    class TokenProcessor(_TokenProcessor):
+        """Interface of the token processor part of RNN-T pipeline"""
+
+    _rnnt_path: str
+    _rnnt_factory_func: Callable[[], RNNT]
+    _global_stats_path: str
+    _sp_model_path: str
+    _right_padding: int
+    _blank: int
+    _sample_rate: int
+    _n_fft: int
+    _n_mels: int
+    _hop_length: int
+    _segment_length: int
+    _right_context_length: int
+
+    def _get_model(self) -> RNNT:
+        model = self._rnnt_factory_func()
+        path = torchaudio.utils._download_asset(self._rnnt_path)
+        state_dict = torch.load(path)
+        model.load_state_dict(state_dict)
+        model.eval()
+        return model
+
+    @property
+    def sample_rate(self) -> int:
+        """Sample rate (in cycles per second) of input waveforms.
+
+        :type: int
+        """
+        return self._sample_rate
+
+    @property
+    def n_fft(self) -> int:
+        """Size of FFT window to use.
+
+        :type: int
+        """
+        return self._n_fft
+
+    @property
+    def n_mels(self) -> int:
+        """Number of mel spectrogram features to extract from input waveforms.
+
+        :type: int
+        """
+        return self._n_mels
+
+    @property
+    def hop_length(self) -> int:
+        """Number of samples between successive frames in input expected by model.
+
+        :type: int
+        """
+        return self._hop_length
+
+    @property
+    def segment_length(self) -> int:
+        """Number of frames in segment in input expected by model.
+
+        :type: int
+        """
+        return self._segment_length
+
+    @property
+    def right_context_length(self) -> int:
+        """Number of frames in right contextual block in input expected by model.
+
+        :type: int
+        """
+        return self._right_context_length
+
+    def get_decoder(self) -> RNNTBeamSearch:
+        """Constructs RNN-T decoder.
+
+        Returns:
+            RNNTBeamSearch
+        """
+        model = self._get_model()
+        return RNNTBeamSearch(model, self._blank)
+
+    def get_feature_extractor(self) -> FeatureExtractor:
+        """Constructs feature extractor for non-streaming (full-context) ASR.
+
+        Returns:
+            FeatureExtractor
+        """
+        local_path = torchaudio.utils._download_asset(self._global_stats_path)
+        return _ModuleFeatureExtractor(
+            torch.nn.Sequential(
+                torchaudio.transforms.MelSpectrogram(
+                    sample_rate=self.sample_rate, n_fft=self.n_fft, n_mels=self.n_mels, hop_length=self.hop_length
+                ),
+                _FunctionalModule(lambda x: x.transpose(1, 0)),
+                _FunctionalModule(lambda x: _piecewise_linear_log(x * _gain)),
+                _GlobalStatsNormalization(local_path),
+                _FunctionalModule(lambda x: torch.nn.functional.pad(x, (0, 0, 0, self._right_padding))),
+            )
+        )
+
+    def get_streaming_feature_extractor(self) -> FeatureExtractor:
+        """Constructs feature extractor for streaming (simultaneous) ASR.
+
+        Returns:
+            FeatureExtractor
+        """
+        local_path = torchaudio.utils._download_asset(self._global_stats_path)
+        return _ModuleFeatureExtractor(
+            torch.nn.Sequential(
+                torchaudio.transforms.MelSpectrogram(
+                    sample_rate=self.sample_rate, n_fft=self.n_fft, n_mels=self.n_mels, hop_length=self.hop_length
+                ),
+                _FunctionalModule(lambda x: x.transpose(1, 0)),
+                _FunctionalModule(lambda x: _piecewise_linear_log(x * _gain)),
+                _GlobalStatsNormalization(local_path),
+            )
+        )
+
+    def get_token_processor(self) -> TokenProcessor:
+        """Constructs token processor.
+
+        Returns:
+            TokenProcessor
+        """
+        local_path = torchaudio.utils._download_asset(self._sp_model_path)
+        return _SentencePieceTokenProcessor(local_path)
+
+
+EMFORMER_RNNT_BASE_LIBRISPEECH = RNNTBundle(
+    _rnnt_path="models/emformer_rnnt_base_librispeech.pt",
+    _rnnt_factory_func=partial(emformer_rnnt_base, num_symbols=4097),
+    _global_stats_path="pipeline-assets/global_stats_rnnt_librispeech.json",
+    _sp_model_path="pipeline-assets/spm_bpe_4096_librispeech.model",
+    _right_padding=4,
+    _blank=4096,
+    _sample_rate=16000,
+    _n_fft=400,
+    _n_mels=80,
+    _hop_length=160,
+    _segment_length=16,
+    _right_context_length=4,
+)
+EMFORMER_RNNT_BASE_LIBRISPEECH.__doc__ = """ASR pipeline based on Emformer-RNNT,
+pretrained on *LibriSpeech* dataset :cite:`7178964`,
+capable of performing both streaming and non-streaming inference.
+
+The underlying model is constructed by :py:func:`torchaudio.models.emformer_rnnt_base`
+and utilizes weights trained on LibriSpeech using training script ``train.py``
+`here <https://github.com/pytorch/audio/tree/main/examples/asr/emformer_rnnt>`__ with default arguments.
+
+Please refer to :py:class:`RNNTBundle` for usage instructions.
+"""

torchaudio/transforms/__init__.py

@@ -0,0 +1,78 @@
+from torchaudio._internal.module_utils import dropping_class_support
+
+from ._multi_channel import MVDR, PSD, RTFMVDR, SoudenMVDR
+from ._transforms import (
+    AddNoise,
+    AmplitudeToDB,
+    ComputeDeltas,
+    Convolve,
+    Deemphasis,
+    Fade,
+    FFTConvolve,
+    FrequencyMasking,
+    GriffinLim,
+    InverseMelScale,
+    InverseSpectrogram,
+    LFCC,
+    Loudness,
+    MelScale,
+    MelSpectrogram,
+    MFCC,
+    MuLawDecoding,
+    MuLawEncoding,
+    PitchShift,
+    Preemphasis,
+    Resample,
+    RNNTLoss as _RNNTLoss,
+    SlidingWindowCmn,
+    SpecAugment,
+    SpectralCentroid,
+    Spectrogram,
+    Speed,
+    SpeedPerturbation,
+    TimeMasking,
+    TimeStretch,
+    Vad,
+    Vol,
+)
+
+RNNTLoss = dropping_class_support(_RNNTLoss)
+
+__all__ = [
+    "AddNoise",
+    "AmplitudeToDB",
+    "ComputeDeltas",
+    "Convolve",
+    "Deemphasis",
+    "Fade",
+    "FFTConvolve",
+    "FrequencyMasking",
+    "GriffinLim",
+    "InverseMelScale",
+    "InverseSpectrogram",
+    "LFCC",
+    "Loudness",
+    "MFCC",
+    "MVDR",
+    "MelScale",
+    "MelSpectrogram",
+    "MuLawDecoding",
+    "MuLawEncoding",
+    "PSD",
+    "PitchShift",
+    "Preemphasis",
+    "RNNTLoss",
+    "RTFMVDR",
+    "Resample",
+    "SlidingWindowCmn",
+    "SoudenMVDR",
+    "SpecAugment",
+    "SpectralCentroid",
+    "Spectrogram",
+    "Speed",
+    "SpeedPerturbation",
+    "TimeMasking",
+    "TimeStretch",
+    "Vad",
+    "Vol",
+]