wehub-resource-sync
574 lines
20 KiB
Python
574 lines
20 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import logging
|
|
import math
|
|
import random
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from torch import nn
|
|
from torchaudio.functional import melscale_fbanks
|
|
from transformers import AutoFeatureExtractor, AutoProcessor, BatchFeature
|
|
from transformers.feature_extraction_sequence_utils import (
|
|
SequenceFeatureExtractor,
|
|
)
|
|
from transformers.processing_utils import ProcessorMixin
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
CONSTANT = 1e-5
|
|
INF_VAL = 10000.0
|
|
|
|
|
|
class FilterbankFeatures(nn.Module):
|
|
"""Featurizer that converts wavs to Mel Spectrograms.
|
|
See AudioToMelSpectrogramPreprocessor for args.
|
|
"""
|
|
|
|
window: torch.Tensor
|
|
fb: torch.Tensor
|
|
|
|
def __init__(
|
|
self,
|
|
sample_rate=16000,
|
|
n_window_size=320,
|
|
n_window_stride=160,
|
|
window="hann",
|
|
normalize="per_feature",
|
|
n_fft=None,
|
|
preemph=0.97,
|
|
nfilt=64,
|
|
lowfreq=0,
|
|
highfreq=None,
|
|
log=True,
|
|
log_zero_guard_type="add",
|
|
log_zero_guard_value=2**-24,
|
|
dither=CONSTANT,
|
|
pad_to=16,
|
|
max_duration=30,
|
|
frame_splicing=1,
|
|
exact_pad=False,
|
|
pad_value=0,
|
|
mag_power=2.0,
|
|
use_grads=False,
|
|
rng=None,
|
|
nb_augmentation_prob=0.0,
|
|
nb_max_freq=4000,
|
|
mel_norm="slaney",
|
|
stft_exact_pad=False,
|
|
stft_conv=False,
|
|
device="cpu",
|
|
):
|
|
super().__init__()
|
|
if stft_conv or stft_exact_pad:
|
|
logger.warning(
|
|
"Using torch_stft is deprecated and has been removed. "
|
|
"The values have been forcibly set to False for "
|
|
"FilterbankFeatures and AudioToMelSpectrogramPreprocessor. "
|
|
"Please set exact_pad to True as needed."
|
|
)
|
|
if exact_pad and n_window_stride % 2 == 1:
|
|
raise NotImplementedError(
|
|
f"{self} received exact_pad == True, but hop_size was odd. "
|
|
"If audio_length % hop_size == 0, the returned spectrogram "
|
|
"would not be of length audio_length // hop_size. "
|
|
"Please use an even hop_size."
|
|
)
|
|
self.log_zero_guard_value = log_zero_guard_value
|
|
if (
|
|
n_window_size is None
|
|
or n_window_stride is None
|
|
or not isinstance(n_window_size, int)
|
|
or not isinstance(n_window_stride, int)
|
|
or n_window_size <= 0
|
|
or n_window_stride <= 0
|
|
):
|
|
raise ValueError(
|
|
f"{self} got an invalid value for either n_window_size or "
|
|
f"n_window_stride. Both must be positive ints."
|
|
)
|
|
|
|
self.sample_rate = sample_rate
|
|
self.win_length = n_window_size
|
|
self.hop_length = n_window_stride
|
|
self.n_fft = n_fft or 2 ** math.ceil(math.log2(self.win_length))
|
|
self.stft_pad_amount = (
|
|
(self.n_fft - self.hop_length) // 2 if exact_pad else None
|
|
)
|
|
self.exact_pad = exact_pad
|
|
self.sample_rate = sample_rate
|
|
self.max_duration = max_duration
|
|
|
|
if exact_pad:
|
|
logger.info("STFT using exact pad")
|
|
torch_windows = {
|
|
"hann": torch.hann_window,
|
|
"hamming": torch.hamming_window,
|
|
"blackman": torch.blackman_window,
|
|
"bartlett": torch.bartlett_window,
|
|
"none": None,
|
|
}
|
|
window_fn = torch_windows.get(window)
|
|
window_tensor = (
|
|
window_fn(self.win_length, periodic=False) if window_fn else None
|
|
)
|
|
self.register_buffer("window", window_tensor)
|
|
|
|
self.normalize = normalize
|
|
self.log = log
|
|
self.dither = dither
|
|
self.frame_splicing = frame_splicing
|
|
self.nfilt = nfilt
|
|
self.preemph = preemph
|
|
self.pad_to = pad_to
|
|
highfreq = highfreq or sample_rate / 2
|
|
self.sample_rate = sample_rate
|
|
# disable pad min duration
|
|
# self.pad_min_duration = 1.0
|
|
self.pad_min_duration = 0.0
|
|
self.pad_direction = "both"
|
|
|
|
filterbanks = melscale_fbanks(
|
|
n_freqs=self.n_fft // 2 + 1,
|
|
f_min=lowfreq,
|
|
f_max=highfreq,
|
|
n_mels=nfilt,
|
|
sample_rate=sample_rate,
|
|
norm=mel_norm,
|
|
mel_scale="slaney",
|
|
).T.unsqueeze(0)
|
|
self.register_buffer("fb", filterbanks)
|
|
|
|
# Calculate maximum sequence length
|
|
max_length = self.get_seq_len(
|
|
torch.tensor(max_duration * sample_rate, dtype=torch.float)
|
|
)
|
|
max_pad = pad_to - (max_length % pad_to) if pad_to > 0 else 0
|
|
self.max_length = max_length + max_pad
|
|
self.pad_value = pad_value
|
|
self.mag_power = mag_power
|
|
|
|
# We want to avoid taking the log of zero
|
|
# There are two options: either adding or clamping to a small value
|
|
if log_zero_guard_type not in ["add", "clamp"]:
|
|
raise ValueError(
|
|
f"{self} received {log_zero_guard_type} for the "
|
|
f"log_zero_guard_type parameter. It must be either 'add' or "
|
|
f"'clamp'."
|
|
)
|
|
|
|
self.use_grads = use_grads
|
|
if not use_grads:
|
|
self.forward = torch.no_grad()(self.forward)
|
|
self._rng = random.Random() if rng is None else rng
|
|
self.nb_augmentation_prob = nb_augmentation_prob
|
|
if self.nb_augmentation_prob > 0.0:
|
|
if nb_max_freq >= sample_rate / 2:
|
|
self.nb_augmentation_prob = 0.0
|
|
else:
|
|
self._nb_max_fft_bin = int((nb_max_freq / sample_rate) * n_fft)
|
|
|
|
# log_zero_guard_value is the the small we want to use, we support
|
|
# an actual number, or "tiny", or "eps"
|
|
self.log_zero_guard_type = log_zero_guard_type
|
|
|
|
assert self.window is not None
|
|
assert self.fb is not None
|
|
self.window = self.window.to(dtype=torch.bfloat16)
|
|
self.fb = self.fb.to(dtype=torch.bfloat16)
|
|
|
|
self.generator = torch.Generator(device=device)
|
|
self.generator.manual_seed(0)
|
|
|
|
@torch._dynamo.disable
|
|
def stft(self, x):
|
|
# disable autocast to get full range of stft values
|
|
with torch.amp.autocast(x.device.type, enabled=False):
|
|
return torch.stft(
|
|
x,
|
|
n_fft=self.n_fft,
|
|
hop_length=self.hop_length,
|
|
win_length=self.win_length,
|
|
center=not self.exact_pad,
|
|
window=self.window.to(dtype=torch.float, device=x.device),
|
|
return_complex=True,
|
|
pad_mode="constant",
|
|
)
|
|
|
|
def log_zero_guard_value_fn(self, x):
|
|
if isinstance(self.log_zero_guard_value, str):
|
|
if self.log_zero_guard_value == "tiny":
|
|
return torch.finfo(x.dtype).tiny
|
|
elif self.log_zero_guard_value == "eps":
|
|
return torch.finfo(x.dtype).eps
|
|
else:
|
|
raise ValueError(
|
|
f"{self} received {self.log_zero_guard_value} for the "
|
|
f"log_zero_guard_type parameter. It must be either a "
|
|
f"number, 'tiny', or 'eps'"
|
|
)
|
|
else:
|
|
return self.log_zero_guard_value
|
|
|
|
def get_seq_len(self, seq_len):
|
|
# Assuming that center is True is stft_pad_amount = 0
|
|
pad_amount = (
|
|
self.stft_pad_amount * 2
|
|
if self.stft_pad_amount is not None
|
|
else self.n_fft // 2 * 2
|
|
)
|
|
seq_len = torch.floor_divide(
|
|
(seq_len + pad_amount - self.n_fft), self.hop_length
|
|
)
|
|
return seq_len.to(dtype=torch.long)
|
|
|
|
@property
|
|
def filter_banks(self):
|
|
return self.fb
|
|
|
|
def splice_frames(self, x, frame_splicing):
|
|
"""Stacks frames together across feature dim
|
|
|
|
input is batch_size, feature_dim, num_frames
|
|
output is batch_size, feature_dim*frame_splicing, num_frames
|
|
|
|
"""
|
|
seq = [x]
|
|
for n in range(1, frame_splicing):
|
|
seq.append(torch.cat([x[:, :, :n], x[:, :, n:]], dim=2))
|
|
return torch.cat(seq, dim=1)
|
|
|
|
def normalize_batch(self, x, seq_len, normalize_type):
|
|
x_mean = None
|
|
x_std = None
|
|
if normalize_type == "per_feature":
|
|
batch_size = x.shape[0]
|
|
max_time = x.shape[2]
|
|
|
|
# When doing stream capture to a graph, item() is not allowed
|
|
# because it calls cudaStreamSynchronize(). Therefore, we are
|
|
# sacrificing some error checking when running with cuda graphs.
|
|
# if (
|
|
# torch.cuda.is_available()
|
|
# and not torch.cuda.is_current_stream_capturing()
|
|
# and torch.any(seq_len == 1).item()
|
|
# ):
|
|
# raise ValueError(
|
|
# "normalize_batch with `per_feature` normalize_type "
|
|
# "received a tensor of length 1. This will result in "
|
|
# "torch.std() returning nan. Make sure your audio length "
|
|
# "has enough samples for a single feature (ex. at least "
|
|
# "`hop_length` for Mel Spectrograms)."
|
|
# )
|
|
time_steps = (
|
|
torch.arange(max_time, device=x.device)
|
|
.unsqueeze(0)
|
|
.expand(batch_size, max_time)
|
|
)
|
|
valid_mask = time_steps < seq_len.unsqueeze(1)
|
|
x_mean_numerator = torch.where(valid_mask.unsqueeze(1), x, 0.0).sum(axis=2)
|
|
x_mean_denominator = valid_mask.sum(axis=1)
|
|
x_mean = x_mean_numerator / x_mean_denominator.unsqueeze(1)
|
|
|
|
# Subtract 1 in the denominator to correct for the bias.
|
|
x_std = torch.sqrt(
|
|
torch.sum(
|
|
torch.where(valid_mask.unsqueeze(1), x - x_mean.unsqueeze(2), 0.0)
|
|
** 2,
|
|
axis=2,
|
|
)
|
|
/ (x_mean_denominator.unsqueeze(1) - 1.0)
|
|
)
|
|
x_std = x_std.masked_fill(
|
|
x_std.isnan(), 0.0
|
|
) # edge case: only 1 frame in denominator
|
|
# make sure x_std is not zero
|
|
x_std += CONSTANT
|
|
return (x - x_mean.unsqueeze(2)) / x_std.unsqueeze(2), x_mean, x_std
|
|
elif normalize_type == "all_features":
|
|
x_mean = torch.zeros(seq_len.shape, dtype=x.dtype, device=x.device)
|
|
x_std = torch.zeros(seq_len.shape, dtype=x.dtype, device=x.device)
|
|
for i in range(x.shape[0]):
|
|
x_mean[i] = x[i, :, : seq_len[i].item()].mean()
|
|
x_std[i] = x[i, :, : seq_len[i].item()].std()
|
|
# make sure x_std is not zero
|
|
x_std += CONSTANT
|
|
return (x - x_mean.view(-1, 1, 1)) / x_std.view(-1, 1, 1), x_mean, x_std
|
|
elif "fixed_mean" in normalize_type and "fixed_std" in normalize_type:
|
|
x_mean = torch.tensor(normalize_type["fixed_mean"], device=x.device)
|
|
x_std = torch.tensor(normalize_type["fixed_std"], device=x.device)
|
|
return (
|
|
(x - x_mean.view(x.shape[0], x.shape[1]).unsqueeze(2))
|
|
/ x_std.view(x.shape[0], x.shape[1]).unsqueeze(2),
|
|
x_mean,
|
|
x_std,
|
|
)
|
|
else:
|
|
return x, x_mean, x_std
|
|
|
|
@torch.compile
|
|
def forward(self, x, seq_len, linear_spec=False):
|
|
if x.shape[1] < self.sample_rate * self.pad_min_duration:
|
|
pad_amount = int(self.sample_rate * self.pad_min_duration) - x.shape[1]
|
|
if self.pad_direction == "right":
|
|
x = F.pad(x, (0, pad_amount), value=self.pad_value)
|
|
elif self.pad_direction == "left":
|
|
x = F.pad(x, (pad_amount, 0), value=self.pad_value)
|
|
elif self.pad_direction == "both":
|
|
left_pad = pad_amount // 2
|
|
right_pad = pad_amount - left_pad
|
|
x = F.pad(x, (left_pad, right_pad), value=self.pad_value)
|
|
else:
|
|
raise ValueError(
|
|
f"{self} received an invalid pad_direction: {self.pad_direction}. "
|
|
f"It must be one of 'left', 'right', or 'both'."
|
|
)
|
|
seq_len = torch.tensor([x.shape[1]], dtype=torch.float, device=x.device)
|
|
|
|
seq_len_time = seq_len
|
|
seq_len_unfixed = self.get_seq_len(seq_len)
|
|
|
|
# fix for seq_len = 0 for streaming; if size was 0, it is always padded
|
|
# to 1, and normalizer fails
|
|
seq_len = torch.where(
|
|
seq_len == 0, torch.zeros_like(seq_len_unfixed), seq_len_unfixed
|
|
)
|
|
|
|
if self.stft_pad_amount is not None:
|
|
x = torch.nn.functional.pad(
|
|
x.unsqueeze(1), (self.stft_pad_amount, self.stft_pad_amount), "constant"
|
|
).squeeze(1)
|
|
|
|
# use dither for inference as well
|
|
if self.dither > 0:
|
|
x += self.dither * torch.randn(
|
|
x.shape, dtype=x.dtype, device=x.device, generator=self.generator
|
|
)
|
|
|
|
# do preemphasis
|
|
if self.preemph is not None:
|
|
timemask = torch.arange(x.shape[1], device=x.device).unsqueeze(
|
|
0
|
|
) < seq_len_time.unsqueeze(1)
|
|
x = torch.cat(
|
|
(x[:, 0].unsqueeze(1), x[:, 1:] - self.preemph * x[:, :-1]), dim=1
|
|
)
|
|
|
|
x = x.masked_fill(~timemask, 0.0)
|
|
|
|
x = self.stft(x)
|
|
|
|
# torch stft returns complex tensor (of shape [B,N,T]); so convert to magnitude
|
|
# guard is needed for sqrt if grads are passed through
|
|
guard = 0 if not self.use_grads else CONSTANT
|
|
x = torch.view_as_real(x)
|
|
x = torch.sqrt(x.pow(2).sum(-1) + guard)
|
|
|
|
# get power spectrum
|
|
if self.mag_power != 1.0:
|
|
x = x.pow(self.mag_power)
|
|
|
|
# return plain spectrogram if required
|
|
if linear_spec:
|
|
return x, seq_len
|
|
|
|
# disable autocast, otherwise it might be automatically casted to fp16
|
|
# on fp16 compatible GPUs and get NaN values for input value of 65520
|
|
with torch.amp.autocast(x.device.type, enabled=False):
|
|
# dot with filterbank energies
|
|
x = torch.matmul(self.fb.to(x.dtype), x)
|
|
|
|
# log features if required
|
|
if self.log:
|
|
if self.log_zero_guard_type == "add":
|
|
x = torch.log(x + self.log_zero_guard_value_fn(x))
|
|
elif self.log_zero_guard_type == "clamp":
|
|
x = torch.log(torch.clamp(x, min=self.log_zero_guard_value_fn(x)))
|
|
else:
|
|
raise ValueError("log_zero_guard_type was not understood")
|
|
|
|
# frame splicing if required
|
|
if self.frame_splicing > 1:
|
|
x = self.splice_frames(x, self.frame_splicing)
|
|
|
|
# normalize if required
|
|
if self.normalize:
|
|
x, _, _ = self.normalize_batch(x, seq_len, normalize_type=self.normalize)
|
|
|
|
# mask to zero any values beyond seq_len in batch, pad to multiple of
|
|
# `pad_to` (for efficiency)
|
|
max_len = x.size(-1)
|
|
mask = torch.arange(max_len, device=x.device)
|
|
mask = mask.repeat(x.size(0), 1) >= seq_len.unsqueeze(1)
|
|
x = x.masked_fill(
|
|
mask.unsqueeze(1).type(torch.bool).to(device=x.device), self.pad_value
|
|
)
|
|
|
|
del mask
|
|
pad_to = self.pad_to
|
|
if pad_to == "max":
|
|
x = nn.functional.pad(
|
|
x, (0, self.max_length - x.size(-1)), value=self.pad_value
|
|
)
|
|
elif pad_to > 0:
|
|
pad_amt = x.size(-1) % pad_to
|
|
if pad_amt != 0:
|
|
x = nn.functional.pad(x, (0, pad_to - pad_amt), value=self.pad_value)
|
|
|
|
return x, seq_len
|
|
|
|
|
|
class CohereASRFeatureExtractor(SequenceFeatureExtractor):
|
|
"""HF-compatible feature extractor wrapping FilterbankFeatures."""
|
|
|
|
model_input_names = ["input_features"]
|
|
|
|
def __init__(
|
|
self,
|
|
feature_size=64,
|
|
sampling_rate=16000,
|
|
padding_value=0.0,
|
|
max_duration=30,
|
|
n_window_size=320,
|
|
n_window_stride=160,
|
|
window="hann",
|
|
normalize="per_feature",
|
|
n_fft=None,
|
|
preemph=0.97,
|
|
lowfreq=0,
|
|
highfreq=None,
|
|
log=True,
|
|
log_zero_guard_type="add",
|
|
log_zero_guard_value=2**-24,
|
|
dither=CONSTANT,
|
|
pad_to=16,
|
|
frame_splicing=1,
|
|
exact_pad=False,
|
|
mag_power=2.0,
|
|
nb_augmentation_prob=0.0,
|
|
nb_max_freq=4000,
|
|
mel_norm="slaney",
|
|
stft_exact_pad=False,
|
|
stft_conv=False,
|
|
device="cpu",
|
|
**kwargs,
|
|
):
|
|
super().__init__(
|
|
feature_size=feature_size,
|
|
sampling_rate=sampling_rate,
|
|
padding_value=padding_value,
|
|
**kwargs,
|
|
)
|
|
self.max_duration = max_duration
|
|
self.hop_length = n_window_stride
|
|
self._device = torch.device(device)
|
|
self._fb_config = dict(
|
|
sample_rate=sampling_rate,
|
|
n_window_size=n_window_size,
|
|
n_window_stride=n_window_stride,
|
|
window=window,
|
|
normalize=normalize,
|
|
n_fft=n_fft,
|
|
preemph=preemph,
|
|
nfilt=feature_size,
|
|
lowfreq=lowfreq,
|
|
highfreq=highfreq,
|
|
log=log,
|
|
log_zero_guard_type=log_zero_guard_type,
|
|
log_zero_guard_value=log_zero_guard_value,
|
|
dither=dither,
|
|
pad_to=pad_to,
|
|
max_duration=max_duration,
|
|
frame_splicing=frame_splicing,
|
|
exact_pad=exact_pad,
|
|
pad_value=padding_value,
|
|
mag_power=mag_power,
|
|
nb_augmentation_prob=nb_augmentation_prob,
|
|
nb_max_freq=nb_max_freq,
|
|
mel_norm=mel_norm,
|
|
stft_exact_pad=stft_exact_pad,
|
|
stft_conv=stft_conv,
|
|
device=device,
|
|
)
|
|
self._filterbank: FilterbankFeatures | None = None
|
|
|
|
@property
|
|
def filterbank(self) -> FilterbankFeatures:
|
|
if self._filterbank is None:
|
|
fb = FilterbankFeatures(**self._fb_config)
|
|
fb.eval()
|
|
self._filterbank = fb.to(self._device)
|
|
return self._filterbank
|
|
|
|
def get_seq_len(self, seq_len):
|
|
return self.filterbank.get_seq_len(seq_len)
|
|
|
|
def __call__(
|
|
self,
|
|
raw_speech,
|
|
sampling_rate=None,
|
|
return_tensors=None,
|
|
**kwargs,
|
|
) -> BatchFeature:
|
|
if isinstance(raw_speech, np.ndarray):
|
|
raw_speech = [raw_speech]
|
|
|
|
seq_len = torch.tensor([s.shape[0] for s in raw_speech])
|
|
|
|
max_len = max(s.shape[0] for s in raw_speech)
|
|
padded = np.zeros((len(raw_speech), max_len), dtype=np.float32)
|
|
for i, s in enumerate(raw_speech):
|
|
padded[i, : s.shape[0]] = s
|
|
|
|
audio_tensor = torch.from_numpy(padded).to(self._device)
|
|
seq_len = seq_len.to(self._device)
|
|
|
|
with torch.no_grad():
|
|
input_features, length = self.filterbank(audio_tensor, seq_len)
|
|
|
|
result = BatchFeature(
|
|
{"input_features": input_features.cpu(), "length": length.cpu()}
|
|
)
|
|
if return_tensors is not None:
|
|
result = result.convert_to_tensors(return_tensors)
|
|
return result
|
|
|
|
|
|
class CohereASRProcessor(ProcessorMixin):
|
|
"""HF-compatible processor combining CohereASRFeatureExtractor and a
|
|
tokenizer."""
|
|
|
|
feature_extractor_class = "CohereASRFeatureExtractor"
|
|
tokenizer_class = "AutoTokenizer"
|
|
|
|
def __init__(self, feature_extractor, tokenizer):
|
|
super().__init__(feature_extractor, tokenizer)
|
|
|
|
def __call__(
|
|
self,
|
|
text=None,
|
|
audio=None,
|
|
sampling_rate=None,
|
|
return_tensors=None,
|
|
**kwargs,
|
|
):
|
|
if audio is not None:
|
|
result = self.feature_extractor(
|
|
audio,
|
|
sampling_rate=sampling_rate,
|
|
return_tensors=return_tensors,
|
|
)
|
|
else:
|
|
result = BatchFeature()
|
|
|
|
if text is not None:
|
|
text_inputs = self.tokenizer(text, return_tensors=return_tensors, **kwargs)
|
|
result["input_ids"] = text_inputs["input_ids"]
|
|
|
|
return result
|
|
|
|
|
|
AutoFeatureExtractor.register("CohereASRFeatureExtractor", CohereASRFeatureExtractor)
|
|
AutoProcessor.register("CohereASRProcessor", CohereASRProcessor)
|