| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from typing import TYPE_CHECKING |
| 5 | |
| 6 | import numpy as np |
| 7 | import torch |
| 8 | import torch.nn.functional as F |
| 9 | from transformers import ( |
| 10 | AutoFeatureExtractor, |
| 11 | BatchFeature, |
| 12 | ) |
| 13 | from transformers.feature_extraction_sequence_utils import SequenceFeatureExtractor |
| 14 | from transformers.processing_utils import ProcessorMixin |
| 15 | from transformers.utils import TensorType |
| 16 | |
| 17 | from vllm.logger import init_logger |
| 18 | from vllm.utils.import_utils import LazyLoader |
| 19 | |
| 20 | if TYPE_CHECKING: |
| 21 | import kaldi_native_fbank as knf |
| 22 | else: |
| 23 | knf = LazyLoader("knf", globals(), "kaldi_native_fbank") |
| 24 | |
| 25 | |
| 26 | logger = init_logger(__name__) |
| 27 | |
| 28 | |
| 29 | class CMVN: |
| 30 | def __init__(self, dim, means, inverse_std_variences): |
| 31 | self.dim, self.means, self.inverse_std_variences = ( |
| 32 | dim, |
| 33 | np.array(means), |
| 34 | np.array(inverse_std_variences), |
| 35 | ) |
| 36 | |
| 37 | def __call__(self, x): |
| 38 | assert x.shape[-1] == self.dim, "CMVN dim mismatch" |
| 39 | out = x - self.means |
| 40 | out = out * self.inverse_std_variences |
| 41 | return out |
| 42 | |
| 43 | |
| 44 | class KaldifeatFbank: |
| 45 | def __init__(self, num_mel_bins=80, frame_length=25, frame_shift=10, dither=1.0): |
| 46 | self.dither = dither |
| 47 | opts = knf.FbankOptions() |
| 48 | opts.frame_opts.dither = dither |
| 49 | opts.mel_opts.num_bins = num_mel_bins |
| 50 | opts.frame_opts.snip_edges = True |
| 51 | opts.mel_opts.debug_mel = False |
| 52 | self.opts = opts |
| 53 | |
| 54 | def __call__(self, sample_rate, wav_np, is_train=False): |
| 55 | dither = self.dither if is_train else 0.0 |
| 56 | self.opts.frame_opts.dither = dither |
| 57 | fbank = knf.OnlineFbank(self.opts) |
| 58 | |
| 59 | fbank.accept_waveform(sample_rate, wav_np.tolist()) |
| 60 | feat = [] |
| 61 | for i in range(fbank.num_frames_ready): |
| 62 | feat.append(fbank.get_frame(i)) |
| 63 | if len(feat) == 0: |
| 64 | print("Check data, len(feat) == 0", wav_np, flush=True) |
| 65 | return np.zeros((0, self.opts.mel_opts.num_bins)) |
| 66 | feat = np.vstack(feat) |
| 67 | return feat |
| 68 | |
| 69 | |
| 70 | class FireRedASR2FeatureExtractor(SequenceFeatureExtractor): |
| 71 | r"""Constructs a FireRedASR2 feature extractor. |
| 72 | |
| 73 | This feature extractor inherits from [`~feature_extraction_sequence_ |
| 74 | utils.SequenceFeatureExtractor`] which contains most of the main |
| 75 | methods. Users should refer to this superclass for more information |
| 76 | regarding those methods. |
| 77 | |
| 78 | This class extracts mel-filter bank features from raw speech using a custom |
| 79 | numpy implementation of the `Short Time Fourier Transform` which should |
| 80 | match pytorch's `torch.stft` equivalent. |
| 81 | |
| 82 | Args: |
| 83 | feature_size (`int`, *optional*, defaults to 80): |
| 84 | The feature dimension of the extracted features. |
| 85 | sampling_rate (`int`, *optional*, defaults to 16000): |
| 86 | The sampling rate at which the audio files should be digitalized |
| 87 | expressed in hertz (Hz). |
| 88 | chunk_length (`int`, *optional*, defaults to 30): |
| 89 | The maximum number of chunks of `sampling_rate` samples used to |
| 90 | trim and pad longer or shorter audio sequences. |
| 91 | padding_value (`float`, *optional*, defaults to 0.0): |
| 92 | Padding value used to pad the audio. Should correspond to silences. |
| 93 | dither (`float`, *optional*, defaults to 0.0): |
| 94 | Adds dithering. In other words, adds a small Gaussian noise to each frame. |
| 95 | E.g. use 0.0001 to add dithering with a normal distribution centered |
| 96 | around 0.0 with standard deviation 0.0001 (assuming [-1,+1] range |
| 97 | of raw_speech). The value 0.0 means no dithering. |
| 98 | Dithering has similar effect as `spectrogram(mel_floor=...)`. It reduces |
| 99 | the high log_mel_fbank values for signals with hard-zero sections, |
| 100 | when VAD cutoff is present in the signal. |
| 101 | |
| 102 | """ |
| 103 | |
| 104 | model_input_names = ["input_features"] |
| 105 | |
| 106 | def __init__( |
| 107 | self, |
| 108 | feature_size=80, |
| 109 | sampling_rate=16000, |
| 110 | chunk_length=30, |
| 111 | padding_value=0.0, |
| 112 | return_attention_mask=False, |
| 113 | dim=80, |
| 114 | means=None, |
| 115 | inverse_std_variences=None, |
| 116 | num_mel_bins=80, |
| 117 | frame_length=25, |
| 118 | frame_shift=10, |
| 119 | dither=0.0, |
| 120 | max_length=3000, |
| 121 | downsample_rate=2, |
| 122 | left_context=3, |
| 123 | right_context=3, |
| 124 | **kwargs, |
| 125 | ): |
| 126 | super().__init__( |
| 127 | feature_size=feature_size, |
| 128 | sampling_rate=sampling_rate, |
| 129 | padding_value=padding_value, |
| 130 | return_attention_mask=return_attention_mask, |
| 131 | **kwargs, |
| 132 | ) |
| 133 | self.chunk_length = chunk_length |
| 134 | self.max_length = max_length |
| 135 | self.dim = dim |
| 136 | self.means = means |
| 137 | self.inverse_std_variences = inverse_std_variences |
| 138 | self.num_mel_bins = num_mel_bins |
| 139 | self.frame_length = frame_length |
| 140 | self.frame_shift = frame_shift |
| 141 | self.dither = dither |
| 142 | self.sampling_rate = sampling_rate |
| 143 | self.downsample_rate = downsample_rate |
| 144 | self.context = left_context + 1 + right_context |
| 145 | |
| 146 | def __call__( |
| 147 | self, |
| 148 | raw_speech: np.ndarray | list[float] | list[np.ndarray] | list[list[float]], |
| 149 | truncation: bool = True, |
| 150 | pad_to_multiple_of: int | None = None, |
| 151 | return_tensors: str | TensorType | None = None, |
| 152 | return_attention_mask: bool | None = None, |
| 153 | padding: str | None = "max_length", |
| 154 | max_length: int | None = None, |
| 155 | sampling_rate: int | None = None, |
| 156 | do_normalize: bool | None = None, |
| 157 | **kwargs, |
| 158 | ) -> BatchFeature: |
| 159 | if sampling_rate != self.sampling_rate: |
| 160 | raise ValueError( |
| 161 | f"The model corresponding to this feature extractor: " |
| 162 | f"{self.__class__.__name__} was trained using a sampling " |
| 163 | f"rate of {self.sampling_rate}. Please make sure that the " |
| 164 | f"provided `raw_speech` input was sampled with " |
| 165 | f"{self.sampling_rate} and not {sampling_rate}." |
| 166 | ) |
| 167 | |
| 168 | def padding_position_is_0(padded_input, input_lengths): |
| 169 | N, T = padded_input.size()[:2] |
| 170 | mask = torch.ones((N, T)).to(padded_input.device) |
| 171 | for i in range(N): |
| 172 | mask[i, input_lengths[i] :] = 0 |
| 173 | mask = mask.unsqueeze(dim=1) |
| 174 | return mask.to(torch.uint8) |
| 175 | |
| 176 | # initialize the CMVN and Fbank objects |
| 177 | self.cmvn = CMVN(self.dim, self.means, self.inverse_std_variences) |
| 178 | self.fbank = KaldifeatFbank( |
| 179 | num_mel_bins=self.num_mel_bins, |
| 180 | frame_length=self.frame_length, |
| 181 | frame_shift=self.frame_shift, |
| 182 | dither=self.dither, |
| 183 | ) |
| 184 | |
| 185 | feats = [] |
| 186 | speech_lengths = [] |
| 187 | fake_token_lengths = [] |
| 188 | for speech in raw_speech: |
| 189 | """ |
| 190 | We must multiply by 32768 here because FireRedASR2 loads audio data |
| 191 | using kaldiio.load_mat, while vLLM loads audio data using pyav. |
| 192 | """ |
| 193 | speech = speech * 32768 |
| 194 | fbank = self.fbank(sampling_rate, speech) |
| 195 | fbank = self.cmvn(fbank) |
| 196 | fbank = torch.from_numpy(fbank).float() |
| 197 | length = fbank.size(0) |
| 198 | feats.append(fbank) |
| 199 | speech_lengths.append(length) |
| 200 | padded_input2 = fbank |
| 201 | padded_input2 = F.pad( |
| 202 | padded_input2, (0, 0, 0, self.context - 1), "constant", 0.0 |
| 203 | ) |
| 204 | src_mask = padding_position_is_0( |
| 205 | padded_input2[None, :, :], torch.tensor([length], dtype=torch.int32) |
| 206 | ) |
| 207 | x_mask = src_mask |
| 208 | mask = x_mask[:, :, :-2:2][:, :, :-2:2] |
| 209 | input_lengths = mask[:, -1, :].sum(dim=-1) |
| 210 | input_lengths = input_lengths // self.downsample_rate |
| 211 | fake_token_len = torch.clamp(input_lengths, min=1) |
| 212 | fake_token_lengths.append(fake_token_len) |
| 213 | |
| 214 | feats = torch.stack(feats, dim=0) |
| 215 | batched_speech = self.pad( |
| 216 | BatchFeature({"input_features": feats}), |
| 217 | padding=padding, |
| 218 | max_length=max_length if max_length else self.max_length, |
| 219 | truncation=truncation, |
| 220 | pad_to_multiple_of=pad_to_multiple_of, |
| 221 | return_attention_mask=return_attention_mask or do_normalize, |
| 222 | ) |
| 223 | |
| 224 | if return_tensors is not None: |
| 225 | batched_speech = batched_speech.convert_to_tensors(return_tensors) |
| 226 | |
| 227 | batched_speech["speech_lengths"] = torch.tensor(speech_lengths) |
| 228 | batched_speech["fake_token_lengths"] = torch.concat(fake_token_lengths) |
| 229 | return batched_speech |
| 230 | |
| 231 | |
| 232 | class FireRedASR2Processor(ProcessorMixin): |
| 233 | r"""Constructs a FireRedASR2 processor which wraps a FireRedASR2 feature extractor |
| 234 | and a FireRedASR2 tokenizer into a single processor. |
| 235 | |
| 236 | [`FireRedASR2Processor`] offers all the functionalities of |
| 237 | [`FireRedASR2FeatureExtractor`] and [`Qwen2Tokenizer`]. See the |
| 238 | [`~FireRedASR2Processor.__call__`] and [`~FireRedASR2Processor.decode`] for more |
| 239 | information. |
| 240 | |
| 241 | Args: |
| 242 | feature_extractor (`FireRedASR2FeatureExtractor`): An instance of |
| 243 | [`FireRedASR2FeatureExtractor`]. |
| 244 | The feature extractor is a required input. |
| 245 | tokenizer (`Qwen2Tokenizer`): |
| 246 | An instance of [`Qwen2Tokenizer`]. The tokenizer is a required |
| 247 | input. |
| 248 | |
| 249 | """ |
| 250 | |
| 251 | feature_extractor_class = "FireRedASR2FeatureExtractor" |
| 252 | tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast") |
| 253 | |
| 254 | def __init__( |
| 255 | self, |
| 256 | feature_extractor, |
| 257 | tokenizer, |
| 258 | audio_token="<|AUDIO|>", |
| 259 | ): |
| 260 | super().__init__(feature_extractor, tokenizer) |
| 261 | self.current_processor = self.feature_extractor |
| 262 | self._in_target_context_manager = False |
| 263 | self.audio_token = ( |
| 264 | tokenizer.audio_token if hasattr(tokenizer, "audio_token") else audio_token |
| 265 | ) |
| 266 | self.audio_token_id = tokenizer.convert_tokens_to_ids(self.audio_token) |
| 267 | |
| 268 | def get_decoder_prompt_ids(self, task=None, language=None, no_timestamps=True): |
| 269 | return self.tokenizer.get_decoder_prompt_ids( |
| 270 | task=task, language=language, no_timestamps=no_timestamps |
| 271 | ) |
| 272 | |
| 273 | def __call__(self, *args, **kwargs): |
| 274 | """Forwards the `audio` argument to FireRedASR2FeatureExtractor's |
| 275 | [`~FireRedASR2FeatureExtractor.__call__`] and the `text` argument to |
| 276 | [`~Qwen2Tokenizer.__call__`]. Please refer to the docstring of the |
| 277 | above two methods for more information. |
| 278 | """ |
| 279 | if self._in_target_context_manager: |
| 280 | return self.current_processor(*args, **kwargs) |
| 281 | |
| 282 | audio = kwargs.pop("audio", None) |
| 283 | sampling_rate = kwargs.pop("sampling_rate", None) |
| 284 | text = kwargs.pop("text", None) |
| 285 | if len(args) > 0: |
| 286 | audio = args[0] |
| 287 | args = args[1:] |
| 288 | |
| 289 | if text is None: |
| 290 | raise ValueError("You need to specify `text` input to process.") |
| 291 | elif isinstance(text, str): |
| 292 | text = [text] |
| 293 | elif not isinstance(text, list) and not isinstance(text[0], str): |
| 294 | raise ValueError( |
| 295 | "Invalid input text. Please provide a string, or a list of strings" |
| 296 | ) |
| 297 | |
| 298 | if audio is not None: |
| 299 | # ensure we have as much audios as audio tokens |
| 300 | num_audio_tokens = sum(sample.count(self.audio_token) for sample in text) |
| 301 | num_audios = 1 if type(audio) is np.ndarray else len(audio) |
| 302 | if num_audio_tokens != num_audios: |
| 303 | raise ValueError( |
| 304 | f"Found {num_audio_tokens} {self.audio_token} token{'s' if num_audio_tokens > 1 else ''} in provided text but received {num_audios} audio{'s' if num_audios > 1 else ''}" # noqa: E501 |
| 305 | ) |
| 306 | inputs = self.feature_extractor( |
| 307 | audio, *args, sampling_rate=sampling_rate, **kwargs |
| 308 | ) |
| 309 | |
| 310 | expanded_text = [] |
| 311 | for sample in text: |
| 312 | replace_str = [] |
| 313 | while self.audio_token in sample: |
| 314 | num_audio_tokens = int(inputs["fake_token_lengths"].item()) |
| 315 | |
| 316 | expanded_audio_token = self.audio_token * num_audio_tokens |
| 317 | |
| 318 | replace_str.append(expanded_audio_token) |
| 319 | sample = sample.replace(self.audio_token, "<placeholder>", 1) |
| 320 | |
| 321 | while "<placeholder>" in sample: |
| 322 | sample = sample.replace("<placeholder>", replace_str.pop(0), 1) |
| 323 | expanded_text.append(sample) |
| 324 | text = expanded_text |
| 325 | |
| 326 | if text is not None: |
| 327 | encodings = self.tokenizer(text, **kwargs) |
| 328 | |
| 329 | if text is None: |
| 330 | return inputs |
| 331 | |
| 332 | elif audio is None: |
| 333 | return encodings |
| 334 | else: |
| 335 | inputs["labels"] = encodings["input_ids"] |
| 336 | |
| 337 | return inputs |
| 338 | |
| 339 | def get_prompt_ids(self, text: str, return_tensors="np"): |
| 340 | return self.tokenizer.get_prompt_ids(text, return_tensors=return_tensors) |
| 341 | |
| 342 | |
| 343 | AutoFeatureExtractor.register( |
| 344 | "FireRedASR2FeatureExtractor", FireRedASR2FeatureExtractor |
| 345 | ) |