File size: 5,742 Bytes
dc0059b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 | from pathlib import Path
import numpy as np
import torch
from torchcodec.decoders import AudioDecoder
from transformers.feature_extraction_sequence_utils import SequenceFeatureExtractor
from transformers.feature_extraction_utils import BatchFeature
from transformers.utils import PaddingStrategy, TensorType, logging
logger = logging.get_logger(__name__)
class MSPAudioFeatureExtractor(SequenceFeatureExtractor):
model_input_names = ["input_values", "padding_mask"]
def __init__(
self,
feature_size: int = 1,
sampling_rate: int = 16000,
padding_value: float = 0.0,
return_attention_mask: bool = True,
do_normalize: bool = True,
**kwargs,
):
super().__init__(
feature_size=feature_size,
sampling_rate=sampling_rate,
padding_value=padding_value,
**kwargs,
)
self.return_attention_mask = return_attention_mask
self.do_normalize = do_normalize
@staticmethod
def zero_mean_unit_var_norm(
input_values: list[np.ndarray],
attention_mask: list[np.ndarray] | None,
padding_value: float = 0.0,
) -> list[np.ndarray]:
"""Normalize each sequence to zero mean and unit variance."""
if attention_mask is not None:
attention_mask = np.array(attention_mask, dtype=np.int32)
normed = []
for vec, length in zip(input_values, attention_mask.sum(-1)):
normed_slice = (vec - vec[:length].mean()) / np.sqrt(
vec[:length].var() + 1e-7
)
if length < normed_slice.shape[0]:
normed_slice[length:] = padding_value
normed.append(normed_slice)
else:
normed = [(x - x.mean()) / np.sqrt(x.var() + 1e-7) for x in input_values]
return normed
def _load_audio(
self,
src: str | Path | bytes | torch.Tensor,
start_seconds: float = 0.0,
stop_seconds: float | None = None,
) -> np.ndarray:
"""Load audio waveform from file path or bytes as a 1-D numpy array."""
audio_decoder = AudioDecoder(source=src, sample_rate=self.sampling_rate)
if stop_seconds is None:
stop_seconds = audio_decoder.metadata.duration_seconds_from_header
waveform = audio_decoder.get_samples_played_in_range(
start_seconds, stop_seconds
).data.numpy()
return waveform.squeeze() # shape: (T,)
def __call__(
self,
raw_speech: (
str
| Path
| bytes
| np.ndarray
| list[str]
| list[Path]
| list[bytes]
| list[float]
| list[np.ndarray]
| list[list[float]]
),
padding: bool | str | PaddingStrategy = False,
max_length: int | None = None,
truncation: bool = False,
pad_to_multiple_of: int | None = None,
return_attention_mask: bool | None = None,
return_tensors: str | TensorType | None = None,
sampling_rate: int | None = None,
**kwargs,
) -> BatchFeature:
"""
Featurize and pad one or several audio sequences.
"""
if sampling_rate is not None and sampling_rate != self.sampling_rate:
raise ValueError(
f"Sampling rate mismatch: expected {self.sampling_rate}, "
f"got {sampling_rate}."
)
is_batched_numpy = isinstance(raw_speech, np.ndarray) and raw_speech.ndim > 1
if is_batched_numpy and raw_speech.ndim > 2:
raise ValueError("Only mono-channel audio is supported.")
is_batched = is_batched_numpy or (
isinstance(raw_speech, (list, tuple))
and isinstance(raw_speech[0], (str, Path, bytes, np.ndarray, list, tuple))
)
if not is_batched:
raw_speech = [raw_speech]
# Load from file paths or bytes
if isinstance(raw_speech[0], (str, Path, bytes)):
raw_speech = [self._load_audio(src) for src in raw_speech]
encoded = BatchFeature({"input_values": raw_speech})
padded = self.pad(
encoded,
padding=padding,
max_length=max_length,
truncation=truncation,
pad_to_multiple_of=pad_to_multiple_of,
return_attention_mask=return_attention_mask,
)
# Ensure float32
vals = padded["input_values"]
if not isinstance(vals[0], np.ndarray):
padded["input_values"] = [np.asarray(a, dtype=np.float32) for a in vals]
elif isinstance(vals[0], np.ndarray) and vals[0].dtype == np.float64:
padded["input_values"] = [a.astype(np.float32) for a in vals]
# Normalize
attn = padded.get("attention_mask")
if attn is not None:
padded["attention_mask"] = [np.asarray(a, dtype=np.int32) for a in attn]
if self.do_normalize:
norm_attn = (
attn
if self._get_padding_strategies(padding, max_length=max_length)
is not PaddingStrategy.DO_NOT_PAD
else None
)
padded["input_values"] = self.zero_mean_unit_var_norm(
padded["input_values"],
attention_mask=norm_attn,
padding_value=self.padding_value,
)
# Rename attention_mask -> padding_mask
if "attention_mask" in padded:
padded["padding_mask"] = padded.pop("attention_mask")
if return_tensors is not None:
padded = padded.convert_to_tensors(return_tensors)
return padded
|