torchaudio 2.9.1__cp311-cp311-manylinux_2_28_aarch64.whl

torchaudio/functional/__init__.py

@@ -0,0 +1,130 @@
+from torchaudio._internal.module_utils import dropping_support
+
+from ._alignment import forced_align as _forced_align, merge_tokens, TokenSpan
+from .filtering import (
+    allpass_biquad,
+    band_biquad,
+    bandpass_biquad,
+    bandreject_biquad,
+    bass_biquad,
+    biquad,
+    contrast,
+    dcshift,
+    deemph_biquad,
+    dither,
+    equalizer_biquad,
+    filtfilt,
+    flanger,
+    gain,
+    highpass_biquad,
+    lfilter,
+    lowpass_biquad,
+    overdrive,
+    phaser,
+    riaa_biquad,
+    treble_biquad,
+    vad,
+)
+
+forced_align = dropping_support(_forced_align)
+
+from .functional import (
+    add_noise,
+    amplitude_to_DB,
+    apply_beamforming,
+    compute_deltas,
+    convolve,
+    create_dct,
+    DB_to_amplitude,
+    deemphasis,
+    detect_pitch_frequency,
+    edit_distance,
+    fftconvolve,
+    frechet_distance,
+    griffinlim,
+    inverse_spectrogram,
+    linear_fbanks,
+    loudness,
+    mask_along_axis,
+    mask_along_axis_iid,
+    melscale_fbanks,
+    mu_law_decoding,
+    mu_law_encoding,
+    mvdr_weights_rtf,
+    mvdr_weights_souden,
+    phase_vocoder,
+    pitch_shift,
+    preemphasis,
+    psd,
+    resample,
+    rnnt_loss,
+    rtf_evd,
+    rtf_power,
+    sliding_window_cmn,
+    spectral_centroid,
+    spectrogram,
+    speed,
+)
+
+__all__ = [
+    "amplitude_to_DB",
+    "compute_deltas",
+    "create_dct",
+    "melscale_fbanks",
+    "linear_fbanks",
+    "DB_to_amplitude",
+    "loudness",
+    "detect_pitch_frequency",
+    "griffinlim",
+    "mask_along_axis",
+    "mask_along_axis_iid",
+    "mu_law_encoding",
+    "mu_law_decoding",
+    "phase_vocoder",
+    "sliding_window_cmn",
+    "spectrogram",
+    "inverse_spectrogram",
+    "spectral_centroid",
+    "allpass_biquad",
+    "band_biquad",
+    "bandpass_biquad",
+    "bandreject_biquad",
+    "bass_biquad",
+    "biquad",
+    "contrast",
+    "dither",
+    "dcshift",
+    "deemph_biquad",
+    "equalizer_biquad",
+    "filtfilt",
+    "flanger",
+    "forced_align",
+    "merge_tokens",
+    "TokenSpan",
+    "gain",
+    "highpass_biquad",
+    "lfilter",
+    "lowpass_biquad",
+    "overdrive",
+    "phaser",
+    "riaa_biquad",
+    "treble_biquad",
+    "vad",
+    "resample",
+    "edit_distance",
+    "pitch_shift",
+    "rnnt_loss",
+    "psd",
+    "mvdr_weights_souden",
+    "mvdr_weights_rtf",
+    "rtf_evd",
+    "rtf_power",
+    "apply_beamforming",
+    "fftconvolve",
+    "convolve",
+    "add_noise",
+    "speed",
+    "preemphasis",
+    "deemphasis",
+    "frechet_distance",
+]

torchaudio/functional/_alignment.py

@@ -0,0 +1,128 @@
+from dataclasses import dataclass
+from typing import List, Optional, Tuple
+
+import torch
+from torch import Tensor
+from torchaudio._extension import fail_if_no_align
+
+__all__ = []
+
+
+@fail_if_no_align
+def forced_align(
+    log_probs: Tensor,
+    targets: Tensor,
+    input_lengths: Optional[Tensor] = None,
+    target_lengths: Optional[Tensor] = None,
+    blank: int = 0,
+) -> Tuple[Tensor, Tensor]:
+    r"""Align a CTC label sequence to an emission.
+
+    .. devices:: CPU CUDA
+
+    .. properties:: TorchScript
+
+    Args:
+        log_probs (Tensor): log probability of CTC emission output.
+            Tensor of shape `(B, T, C)`. where `B` is the batch size, `T` is the input length,
+            `C` is the number of characters in alphabet including blank.
+        targets (Tensor): Target sequence. Tensor of shape `(B, L)`,
+            where `L` is the target length.
+        input_lengths (Tensor or None, optional):
+            Lengths of the inputs (max value must each be <= `T`). 1-D Tensor of shape `(B,)`.
+        target_lengths (Tensor or None, optional):
+            Lengths of the targets. 1-D Tensor of shape `(B,)`.
+        blank_id (int, optional): The index of blank symbol in CTC emission. (Default: 0)
+
+    Returns:
+        Tuple(Tensor, Tensor):
+            Tensor: Label for each time step in the alignment path computed using forced alignment.
+
+            Tensor: Log probability scores of the labels for each time step.
+
+    Note:
+        The sequence length of `log_probs` must satisfy:
+
+
+        .. math::
+            L_{\text{log\_probs}} \ge L_{\text{label}} + N_{\text{repeat}}
+
+        where :math:`N_{\text{repeat}}` is the number of consecutively repeated tokens.
+        For example, in str `"aabbc"`, the number of repeats are `2`.
+
+    Note:
+        The current version only supports ``batch_size==1``.
+    """
+    if blank in targets:
+        raise ValueError(f"targets Tensor shouldn't contain blank index. Found {targets}.")
+    if torch.max(targets) >= log_probs.shape[-1]:
+        raise ValueError("targets values must be less than the CTC dimension")
+
+    if input_lengths is None:
+        batch_size, length = log_probs.size(0), log_probs.size(1)
+        input_lengths = torch.full((batch_size,), length, dtype=torch.int64, device=log_probs.device)
+    if target_lengths is None:
+        batch_size, length = targets.size(0), targets.size(1)
+        target_lengths = torch.full((batch_size,), length, dtype=torch.int64, device=targets.device)
+
+    # For TorchScript compatibility
+    assert input_lengths is not None
+    assert target_lengths is not None
+
+    paths, scores = torch.ops.torchaudio.forced_align(log_probs, targets, input_lengths, target_lengths, blank)
+    return paths, scores[:, torch.arange(scores.shape[1]), paths[0]]
+
+
+@dataclass
+class TokenSpan:
+    """TokenSpan()
+    Token with time stamps and score. Returned by :py:func:`merge_tokens`.
+    """
+
+    token: int
+    """The token"""
+    start: int
+    """The start time (inclusive) in emission time axis."""
+    end: int
+    """The end time (exclusive) in emission time axis."""
+    score: float
+    """The score of the this token."""
+
+    def __len__(self) -> int:
+        """Returns the time span"""
+        return self.end - self.start
+
+
+def merge_tokens(tokens: Tensor, scores: Tensor, blank: int = 0) -> List[TokenSpan]:
+    """Removes repeated tokens and blank tokens from the given CTC token sequence.
+
+    Args:
+        tokens (Tensor): Alignment tokens (unbatched) returned from :py:func:`forced_align`.
+            Shape: `(time, )`.
+        scores (Tensor): Alignment scores (unbatched) returned from :py:func:`forced_align`.
+            Shape: `(time, )`. When computing the token-size score, the given score is averaged
+            across the corresponding time span.
+
+    Returns:
+        list of TokenSpan
+
+    Example:
+        >>> aligned_tokens, scores = forced_align(emission, targets, input_lengths, target_lengths)
+        >>> token_spans = merge_tokens(aligned_tokens[0], scores[0])
+    """
+    if tokens.ndim != 1 or scores.ndim != 1:
+        raise ValueError("`tokens` and `scores` must be 1D Tensor.")
+    if len(tokens) != len(scores):
+        raise ValueError("`tokens` and `scores` must be the same length.")
+
+    diff = torch.diff(
+        tokens, prepend=torch.tensor([-1], device=tokens.device), append=torch.tensor([-1], device=tokens.device)
+    )
+    changes_wo_blank = torch.nonzero((diff != 0)).squeeze().tolist()
+    tokens = tokens.tolist()
+    spans = [
+        TokenSpan(token=token, start=start, end=end, score=scores[start:end].mean().item())
+        for start, end in zip(changes_wo_blank[:-1], changes_wo_blank[1:])
+        if (token := tokens[start]) != blank
+    ]
+    return spans