| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import logging |
| 4 | import math |
| 5 | import random |
| 6 | |
| 7 | import numpy as np |
| 8 | import torch |
| 9 | import torch.nn.functional as F |
| 10 | from torch import nn |
| 11 | from torchaudio.functional import melscale_fbanks |
| 12 | from transformers import AutoFeatureExtractor, AutoProcessor, BatchFeature |
| 13 | from transformers.feature_extraction_sequence_utils import ( |
| 14 | SequenceFeatureExtractor, |
| 15 | ) |
| 16 | from transformers.processing_utils import ProcessorMixin |
| 17 | |
| 18 | logger = logging.getLogger(__name__) |
| 19 | |
| 20 | CONSTANT = 1e-5 |
| 21 | INF_VAL = 10000.0 |
| 22 | |
| 23 | |
| 24 | class FilterbankFeatures(nn.Module): |
| 25 | """Featurizer that converts wavs to Mel Spectrograms. |
| 26 | See AudioToMelSpectrogramPreprocessor for args. |
| 27 | """ |
| 28 | |
| 29 | window: torch.Tensor |
| 30 | fb: torch.Tensor |
| 31 | |
| 32 | def __init__( |
| 33 | self, |
| 34 | sample_rate=16000, |
| 35 | n_window_size=320, |
| 36 | n_window_stride=160, |
| 37 | window="hann", |
| 38 | normalize="per_feature", |
| 39 | n_fft=None, |
| 40 | preemph=0.97, |
| 41 | nfilt=64, |
| 42 | lowfreq=0, |
| 43 | highfreq=None, |
| 44 | log=True, |
| 45 | log_zero_guard_type="add", |
| 46 | log_zero_guard_value=2**-24, |
| 47 | dither=CONSTANT, |
| 48 | pad_to=16, |
| 49 | max_duration=30, |
| 50 | frame_splicing=1, |
| 51 | exact_pad=False, |
| 52 | pad_value=0, |
| 53 | mag_power=2.0, |
| 54 | use_grads=False, |
| 55 | rng=None, |
| 56 | nb_augmentation_prob=0.0, |
| 57 | nb_max_freq=4000, |
| 58 | mel_norm="slaney", |
| 59 | stft_exact_pad=False, |
| 60 | stft_conv=False, |
| 61 | device="cpu", |
| 62 | ): |
| 63 | super().__init__() |
| 64 | if stft_conv or stft_exact_pad: |
| 65 | logger.warning( |
| 66 | "Using torch_stft is deprecated and has been removed. " |
| 67 | "The values have been forcibly set to False for " |
| 68 | "FilterbankFeatures and AudioToMelSpectrogramPreprocessor. " |
| 69 | "Please set exact_pad to True as needed." |
| 70 | ) |
| 71 | if exact_pad and n_window_stride % 2 == 1: |
| 72 | raise NotImplementedError( |
| 73 | f"{self} received exact_pad == True, but hop_size was odd. " |
| 74 | "If audio_length % hop_size == 0, the returned spectrogram " |
| 75 | "would not be of length audio_length // hop_size. " |
| 76 | "Please use an even hop_size." |
| 77 | ) |
| 78 | self.log_zero_guard_value = log_zero_guard_value |
| 79 | if ( |
| 80 | n_window_size is None |
| 81 | or n_window_stride is None |
| 82 | or not isinstance(n_window_size, int) |
| 83 | or not isinstance(n_window_stride, int) |
| 84 | or n_window_size <= 0 |
| 85 | or n_window_stride <= 0 |
| 86 | ): |
| 87 | raise ValueError( |
| 88 | f"{self} got an invalid value for either n_window_size or " |
| 89 | f"n_window_stride. Both must be positive ints." |
| 90 | ) |
| 91 | |
| 92 | self.sample_rate = sample_rate |
| 93 | self.win_length = n_window_size |
| 94 | self.hop_length = n_window_stride |
| 95 | self.n_fft = n_fft or 2 ** math.ceil(math.log2(self.win_length)) |
| 96 | self.stft_pad_amount = ( |
| 97 | (self.n_fft - self.hop_length) // 2 if exact_pad else None |
| 98 | ) |
| 99 | self.exact_pad = exact_pad |
| 100 | self.sample_rate = sample_rate |
| 101 | self.max_duration = max_duration |
| 102 | |
| 103 | if exact_pad: |
| 104 | logger.info("STFT using exact pad") |
| 105 | torch_windows = { |
| 106 | "hann": torch.hann_window, |
| 107 | "hamming": torch.hamming_window, |
| 108 | "blackman": torch.blackman_window, |
| 109 | "bartlett": torch.bartlett_window, |
| 110 | "none": None, |
| 111 | } |
| 112 | window_fn = torch_windows.get(window) |
| 113 | window_tensor = ( |
| 114 | window_fn(self.win_length, periodic=False) if window_fn else None |
| 115 | ) |
| 116 | self.register_buffer("window", window_tensor) |
| 117 | |
| 118 | self.normalize = normalize |
| 119 | self.log = log |
| 120 | self.dither = dither |
| 121 | self.frame_splicing = frame_splicing |
| 122 | self.nfilt = nfilt |
| 123 | self.preemph = preemph |
| 124 | self.pad_to = pad_to |
| 125 | highfreq = highfreq or sample_rate / 2 |
| 126 | self.sample_rate = sample_rate |
| 127 | # disable pad min duration |
| 128 | # self.pad_min_duration = 1.0 |
| 129 | self.pad_min_duration = 0.0 |
| 130 | self.pad_direction = "both" |
| 131 | |
| 132 | filterbanks = melscale_fbanks( |
| 133 | n_freqs=self.n_fft // 2 + 1, |
| 134 | f_min=lowfreq, |
| 135 | f_max=highfreq, |
| 136 | n_mels=nfilt, |
| 137 | sample_rate=sample_rate, |
| 138 | norm=mel_norm, |
| 139 | mel_scale="slaney", |
| 140 | ).T.unsqueeze(0) |
| 141 | self.register_buffer("fb", filterbanks) |
| 142 | |
| 143 | # Calculate maximum sequence length |
| 144 | max_length = self.get_seq_len( |
| 145 | torch.tensor(max_duration * sample_rate, dtype=torch.float) |
| 146 | ) |
| 147 | max_pad = pad_to - (max_length % pad_to) if pad_to > 0 else 0 |
| 148 | self.max_length = max_length + max_pad |
| 149 | self.pad_value = pad_value |
| 150 | self.mag_power = mag_power |
| 151 | |
| 152 | # We want to avoid taking the log of zero |
| 153 | # There are two options: either adding or clamping to a small value |
| 154 | if log_zero_guard_type not in ["add", "clamp"]: |
| 155 | raise ValueError( |
| 156 | f"{self} received {log_zero_guard_type} for the " |
| 157 | f"log_zero_guard_type parameter. It must be either 'add' or " |
| 158 | f"'clamp'." |
| 159 | ) |
| 160 | |
| 161 | self.use_grads = use_grads |
| 162 | if not use_grads: |
| 163 | self.forward = torch.no_grad()(self.forward) |
| 164 | self._rng = random.Random() if rng is None else rng |
| 165 | self.nb_augmentation_prob = nb_augmentation_prob |
| 166 | if self.nb_augmentation_prob > 0.0: |
| 167 | if nb_max_freq >= sample_rate / 2: |
| 168 | self.nb_augmentation_prob = 0.0 |
| 169 | else: |
| 170 | self._nb_max_fft_bin = int((nb_max_freq / sample_rate) * n_fft) |
| 171 | |
| 172 | # log_zero_guard_value is the the small we want to use, we support |
| 173 | # an actual number, or "tiny", or "eps" |
| 174 | self.log_zero_guard_type = log_zero_guard_type |
| 175 | |
| 176 | assert self.window is not None |
| 177 | assert self.fb is not None |
| 178 | self.window = self.window.to(dtype=torch.bfloat16) |
| 179 | self.fb = self.fb.to(dtype=torch.bfloat16) |
| 180 | |
| 181 | self.generator = torch.Generator(device=device) |
| 182 | self.generator.manual_seed(0) |
| 183 | |
| 184 | @torch._dynamo.disable |
| 185 | def stft(self, x): |
| 186 | # disable autocast to get full range of stft values |
| 187 | with torch.amp.autocast(x.device.type, enabled=False): |
| 188 | return torch.stft( |
| 189 | x, |
| 190 | n_fft=self.n_fft, |
| 191 | hop_length=self.hop_length, |
| 192 | win_length=self.win_length, |
| 193 | center=not self.exact_pad, |
| 194 | window=self.window.to(dtype=torch.float, device=x.device), |
| 195 | return_complex=True, |
| 196 | pad_mode="constant", |
| 197 | ) |
| 198 | |
| 199 | def log_zero_guard_value_fn(self, x): |
| 200 | if isinstance(self.log_zero_guard_value, str): |
| 201 | if self.log_zero_guard_value == "tiny": |
| 202 | return torch.finfo(x.dtype).tiny |
| 203 | elif self.log_zero_guard_value == "eps": |
| 204 | return torch.finfo(x.dtype).eps |
| 205 | else: |
| 206 | raise ValueError( |
| 207 | f"{self} received {self.log_zero_guard_value} for the " |
| 208 | f"log_zero_guard_type parameter. It must be either a " |
| 209 | f"number, 'tiny', or 'eps'" |
| 210 | ) |
| 211 | else: |
| 212 | return self.log_zero_guard_value |
| 213 | |
| 214 | def get_seq_len(self, seq_len): |
| 215 | # Assuming that center is True is stft_pad_amount = 0 |
| 216 | pad_amount = ( |
| 217 | self.stft_pad_amount * 2 |
| 218 | if self.stft_pad_amount is not None |
| 219 | else self.n_fft // 2 * 2 |
| 220 | ) |
| 221 | seq_len = torch.floor_divide( |
| 222 | (seq_len + pad_amount - self.n_fft), self.hop_length |
| 223 | ) |
| 224 | return seq_len.to(dtype=torch.long) |
| 225 | |
| 226 | @property |
| 227 | def filter_banks(self): |
| 228 | return self.fb |
| 229 | |
| 230 | def splice_frames(self, x, frame_splicing): |
| 231 | """Stacks frames together across feature dim. |
| 232 | |
| 233 | input is batch_size, feature_dim, num_frames |
| 234 | output is batch_size, feature_dim*frame_splicing, num_frames |
| 235 | |
| 236 | """ |
| 237 | seq = [x] |
| 238 | for n in range(1, frame_splicing): |
| 239 | seq.append(torch.cat([x[:, :, :n], x[:, :, n:]], dim=2)) |
| 240 | return torch.cat(seq, dim=1) |
| 241 | |
| 242 | def normalize_batch(self, x, seq_len, normalize_type): |
| 243 | x_mean = None |
| 244 | x_std = None |
| 245 | if normalize_type == "per_feature": |
| 246 | batch_size = x.shape[0] |
| 247 | max_time = x.shape[2] |
| 248 | |
| 249 | # When doing stream capture to a graph, item() is not allowed |
| 250 | # because it calls cudaStreamSynchronize(). Therefore, we are |
| 251 | # sacrificing some error checking when running with cuda graphs. |
| 252 | # if ( |
| 253 | # torch.cuda.is_available() |
| 254 | # and not torch.cuda.is_current_stream_capturing() |
| 255 | # and torch.any(seq_len == 1).item() |
| 256 | # ): |
| 257 | # raise ValueError( |
| 258 | # "normalize_batch with `per_feature` normalize_type " |
| 259 | # "received a tensor of length 1. This will result in " |
| 260 | # "torch.std() returning nan. Make sure your audio length " |
| 261 | # "has enough samples for a single feature (ex. at least " |
| 262 | # "`hop_length` for Mel Spectrograms)." |
| 263 | # ) |
| 264 | time_steps = ( |
| 265 | torch.arange(max_time, device=x.device) |
| 266 | .unsqueeze(0) |
| 267 | .expand(batch_size, max_time) |
| 268 | ) |
| 269 | valid_mask = time_steps < seq_len.unsqueeze(1) |
| 270 | x_mean_numerator = torch.where(valid_mask.unsqueeze(1), x, 0.0).sum(axis=2) |
| 271 | x_mean_denominator = valid_mask.sum(axis=1) |
| 272 | x_mean = x_mean_numerator / x_mean_denominator.unsqueeze(1) |
| 273 | |
| 274 | # Subtract 1 in the denominator to correct for the bias. |
| 275 | x_std = torch.sqrt( |
| 276 | torch.sum( |
| 277 | torch.where(valid_mask.unsqueeze(1), x - x_mean.unsqueeze(2), 0.0) |
| 278 | ** 2, |
| 279 | axis=2, |
| 280 | ) |
| 281 | / (x_mean_denominator.unsqueeze(1) - 1.0) |
| 282 | ) |
| 283 | x_std = x_std.masked_fill( |
| 284 | x_std.isnan(), 0.0 |
| 285 | ) # edge case: only 1 frame in denominator |
| 286 | # make sure x_std is not zero |
| 287 | x_std += CONSTANT |
| 288 | return (x - x_mean.unsqueeze(2)) / x_std.unsqueeze(2), x_mean, x_std |
| 289 | elif normalize_type == "all_features": |
| 290 | x_mean = torch.zeros(seq_len.shape, dtype=x.dtype, device=x.device) |
| 291 | x_std = torch.zeros(seq_len.shape, dtype=x.dtype, device=x.device) |
| 292 | for i in range(x.shape[0]): |
| 293 | x_mean[i] = x[i, :, : seq_len[i].item()].mean() |
| 294 | x_std[i] = x[i, :, : seq_len[i].item()].std() |
| 295 | # make sure x_std is not zero |
| 296 | x_std += CONSTANT |
| 297 | return (x - x_mean.view(-1, 1, 1)) / x_std.view(-1, 1, 1), x_mean, x_std |
| 298 | elif "fixed_mean" in normalize_type and "fixed_std" in normalize_type: |
| 299 | x_mean = torch.tensor(normalize_type["fixed_mean"], device=x.device) |
| 300 | x_std = torch.tensor(normalize_type["fixed_std"], device=x.device) |
| 301 | return ( |
| 302 | (x - x_mean.view(x.shape[0], x.shape[1]).unsqueeze(2)) |
| 303 | / x_std.view(x.shape[0], x.shape[1]).unsqueeze(2), |
| 304 | x_mean, |
| 305 | x_std, |
| 306 | ) |
| 307 | else: |
| 308 | return x, x_mean, x_std |
| 309 | |
| 310 | @torch.compile |
| 311 | def forward(self, x, seq_len, linear_spec=False): |
| 312 | if x.shape[1] < self.sample_rate * self.pad_min_duration: |
| 313 | pad_amount = int(self.sample_rate * self.pad_min_duration) - x.shape[1] |
| 314 | if self.pad_direction == "right": |
| 315 | x = F.pad(x, (0, pad_amount), value=self.pad_value) |
| 316 | elif self.pad_direction == "left": |
| 317 | x = F.pad(x, (pad_amount, 0), value=self.pad_value) |
| 318 | elif self.pad_direction == "both": |
| 319 | left_pad = pad_amount // 2 |
| 320 | right_pad = pad_amount - left_pad |
| 321 | x = F.pad(x, (left_pad, right_pad), value=self.pad_value) |
| 322 | else: |
| 323 | raise ValueError( |
| 324 | f"{self} received an invalid pad_direction: {self.pad_direction}. " |
| 325 | f"It must be one of 'left', 'right', or 'both'." |
| 326 | ) |
| 327 | seq_len = torch.tensor([x.shape[1]], dtype=torch.float, device=x.device) |
| 328 | |
| 329 | seq_len_time = seq_len |
| 330 | seq_len_unfixed = self.get_seq_len(seq_len) |
| 331 | |
| 332 | # fix for seq_len = 0 for streaming; if size was 0, it is always padded |
| 333 | # to 1, and normalizer fails |
| 334 | seq_len = torch.where( |
| 335 | seq_len == 0, torch.zeros_like(seq_len_unfixed), seq_len_unfixed |
| 336 | ) |
| 337 | |
| 338 | if self.stft_pad_amount is not None: |
| 339 | x = torch.nn.functional.pad( |
| 340 | x.unsqueeze(1), (self.stft_pad_amount, self.stft_pad_amount), "constant" |
| 341 | ).squeeze(1) |
| 342 | |
| 343 | # use dither for inference as well |
| 344 | if self.dither > 0: |
| 345 | x += self.dither * torch.randn( |
| 346 | x.shape, dtype=x.dtype, device=x.device, generator=self.generator |
| 347 | ) |
| 348 | |
| 349 | # do preemphasis |
| 350 | if self.preemph is not None: |
| 351 | timemask = torch.arange(x.shape[1], device=x.device).unsqueeze( |
| 352 | 0 |
| 353 | ) < seq_len_time.unsqueeze(1) |
| 354 | x = torch.cat( |
| 355 | (x[:, 0].unsqueeze(1), x[:, 1:] - self.preemph * x[:, :-1]), dim=1 |
| 356 | ) |
| 357 | |
| 358 | x = x.masked_fill(~timemask, 0.0) |
| 359 | |
| 360 | x = self.stft(x) |
| 361 | |
| 362 | # torch stft returns complex tensor (of shape [B,N,T]); so convert to magnitude |
| 363 | # guard is needed for sqrt if grads are passed through |
| 364 | guard = 0 if not self.use_grads else CONSTANT |
| 365 | x = torch.view_as_real(x) |
| 366 | x = torch.sqrt(x.pow(2).sum(-1) + guard) |
| 367 | |
| 368 | # get power spectrum |
| 369 | if self.mag_power != 1.0: |
| 370 | x = x.pow(self.mag_power) |
| 371 | |
| 372 | # return plain spectrogram if required |
| 373 | if linear_spec: |
| 374 | return x, seq_len |
| 375 | |
| 376 | # disable autocast, otherwise it might be automatically casted to fp16 |
| 377 | # on fp16 compatible GPUs and get NaN values for input value of 65520 |
| 378 | with torch.amp.autocast(x.device.type, enabled=False): |
| 379 | # dot with filterbank energies |
| 380 | x = torch.matmul(self.fb.to(x.dtype), x) |
| 381 | |
| 382 | # log features if required |
| 383 | if self.log: |
| 384 | if self.log_zero_guard_type == "add": |
| 385 | x = torch.log(x + self.log_zero_guard_value_fn(x)) |
| 386 | elif self.log_zero_guard_type == "clamp": |
| 387 | x = torch.log(torch.clamp(x, min=self.log_zero_guard_value_fn(x))) |
| 388 | else: |
| 389 | raise ValueError("log_zero_guard_type was not understood") |
| 390 | |
| 391 | # frame splicing if required |
| 392 | if self.frame_splicing > 1: |
| 393 | x = self.splice_frames(x, self.frame_splicing) |
| 394 | |
| 395 | # normalize if required |
| 396 | if self.normalize: |
| 397 | x, _, _ = self.normalize_batch(x, seq_len, normalize_type=self.normalize) |
| 398 | |
| 399 | # mask to zero any values beyond seq_len in batch, pad to multiple of |
| 400 | # `pad_to` (for efficiency) |
| 401 | max_len = x.size(-1) |
| 402 | mask = torch.arange(max_len, device=x.device) |
| 403 | mask = mask.repeat(x.size(0), 1) >= seq_len.unsqueeze(1) |
| 404 | x = x.masked_fill( |
| 405 | mask.unsqueeze(1).type(torch.bool).to(device=x.device), self.pad_value |
| 406 | ) |
| 407 | |
| 408 | del mask |
| 409 | pad_to = self.pad_to |
| 410 | if pad_to == "max": |
| 411 | x = nn.functional.pad( |
| 412 | x, (0, self.max_length - x.size(-1)), value=self.pad_value |
| 413 | ) |
| 414 | elif pad_to > 0: |
| 415 | pad_amt = x.size(-1) % pad_to |
| 416 | if pad_amt != 0: |
| 417 | x = nn.functional.pad(x, (0, pad_to - pad_amt), value=self.pad_value) |
| 418 | |
| 419 | return x, seq_len |
| 420 | |
| 421 | |
| 422 | class CohereASRFeatureExtractor(SequenceFeatureExtractor): |
| 423 | """HF-compatible feature extractor wrapping FilterbankFeatures.""" |
| 424 | |
| 425 | model_input_names = ["input_features"] |
| 426 | |
| 427 | def __init__( |
| 428 | self, |
| 429 | feature_size=64, |
| 430 | sampling_rate=16000, |
| 431 | padding_value=0.0, |
| 432 | max_duration=30, |
| 433 | n_window_size=320, |
| 434 | n_window_stride=160, |
| 435 | window="hann", |
| 436 | normalize="per_feature", |
| 437 | n_fft=None, |
| 438 | preemph=0.97, |
| 439 | lowfreq=0, |
| 440 | highfreq=None, |
| 441 | log=True, |
| 442 | log_zero_guard_type="add", |
| 443 | log_zero_guard_value=2**-24, |
| 444 | dither=CONSTANT, |
| 445 | pad_to=16, |
| 446 | frame_splicing=1, |
| 447 | exact_pad=False, |
| 448 | mag_power=2.0, |
| 449 | nb_augmentation_prob=0.0, |
| 450 | nb_max_freq=4000, |
| 451 | mel_norm="slaney", |
| 452 | stft_exact_pad=False, |
| 453 | stft_conv=False, |
| 454 | device="cpu", |
| 455 | **kwargs, |
| 456 | ): |
| 457 | super().__init__( |
| 458 | feature_size=feature_size, |
| 459 | sampling_rate=sampling_rate, |
| 460 | padding_value=padding_value, |
| 461 | **kwargs, |
| 462 | ) |
| 463 | self.max_duration = max_duration |
| 464 | self.hop_length = n_window_stride |
| 465 | self._device = torch.device(device) |
| 466 | self._fb_config = dict( |
| 467 | sample_rate=sampling_rate, |
| 468 | n_window_size=n_window_size, |
| 469 | n_window_stride=n_window_stride, |
| 470 | window=window, |
| 471 | normalize=normalize, |
| 472 | n_fft=n_fft, |
| 473 | preemph=preemph, |
| 474 | nfilt=feature_size, |
| 475 | lowfreq=lowfreq, |
| 476 | highfreq=highfreq, |
| 477 | log=log, |
| 478 | log_zero_guard_type=log_zero_guard_type, |
| 479 | log_zero_guard_value=log_zero_guard_value, |
| 480 | dither=dither, |
| 481 | pad_to=pad_to, |
| 482 | max_duration=max_duration, |
| 483 | frame_splicing=frame_splicing, |
| 484 | exact_pad=exact_pad, |
| 485 | pad_value=padding_value, |
| 486 | mag_power=mag_power, |
| 487 | nb_augmentation_prob=nb_augmentation_prob, |
| 488 | nb_max_freq=nb_max_freq, |
| 489 | mel_norm=mel_norm, |
| 490 | stft_exact_pad=stft_exact_pad, |
| 491 | stft_conv=stft_conv, |
| 492 | device=device, |
| 493 | ) |
| 494 | self._filterbank: FilterbankFeatures | None = None |
| 495 | |
| 496 | @property |
| 497 | def filterbank(self) -> FilterbankFeatures: |
| 498 | if self._filterbank is None: |
| 499 | fb = FilterbankFeatures(**self._fb_config) |
| 500 | fb.eval() |
| 501 | self._filterbank = fb.to(self._device) |
| 502 | return self._filterbank |
| 503 | |
| 504 | def get_seq_len(self, seq_len): |
| 505 | return self.filterbank.get_seq_len(seq_len) |
| 506 | |
| 507 | def __call__( |
| 508 | self, |
| 509 | raw_speech, |
| 510 | sampling_rate=None, |
| 511 | return_tensors=None, |
| 512 | **kwargs, |
| 513 | ) -> BatchFeature: |
| 514 | if isinstance(raw_speech, np.ndarray): |
| 515 | raw_speech = [raw_speech] |
| 516 | |
| 517 | seq_len = torch.tensor([s.shape[0] for s in raw_speech]) |
| 518 | |
| 519 | max_len = max(s.shape[0] for s in raw_speech) |
| 520 | padded = np.zeros((len(raw_speech), max_len), dtype=np.float32) |
| 521 | for i, s in enumerate(raw_speech): |
| 522 | padded[i, : s.shape[0]] = s |
| 523 | |
| 524 | audio_tensor = torch.from_numpy(padded).to(self._device) |
| 525 | seq_len = seq_len.to(self._device) |
| 526 | |
| 527 | with torch.no_grad(): |
| 528 | input_features, length = self.filterbank(audio_tensor, seq_len) |
| 529 | |
| 530 | result = BatchFeature( |
| 531 | {"input_features": input_features.cpu(), "length": length.cpu()} |
| 532 | ) |
| 533 | if return_tensors is not None: |
| 534 | result = result.convert_to_tensors(return_tensors) |
| 535 | return result |
| 536 | |
| 537 | |
| 538 | class CohereASRProcessor(ProcessorMixin): |
| 539 | """HF-compatible processor combining CohereASRFeatureExtractor and a |
| 540 | tokenizer.""" |
| 541 | |
| 542 | feature_extractor_class = "CohereASRFeatureExtractor" |
| 543 | tokenizer_class = "AutoTokenizer" |
| 544 | |
| 545 | def __init__(self, feature_extractor, tokenizer): |
| 546 | super().__init__(feature_extractor, tokenizer) |
| 547 | |
| 548 | def __call__( |
| 549 | self, |
| 550 | text=None, |
| 551 | audio=None, |
| 552 | sampling_rate=None, |
| 553 | return_tensors=None, |
| 554 | **kwargs, |
| 555 | ): |
| 556 | if audio is not None: |
| 557 | result = self.feature_extractor( |
| 558 | audio, |
| 559 | sampling_rate=sampling_rate, |
| 560 | return_tensors=return_tensors, |
| 561 | ) |
| 562 | else: |
| 563 | result = BatchFeature() |
| 564 | |
| 565 | if text is not None: |
| 566 | text_inputs = self.tokenizer(text, return_tensors=return_tensors, **kwargs) |
| 567 | result["input_ids"] = text_inputs["input_ids"] |
| 568 | |
| 569 | return result |
| 570 | |
| 571 | |
| 572 | AutoFeatureExtractor.register("CohereASRFeatureExtractor", CohereASRFeatureExtractor) |
| 573 | AutoProcessor.register("CohereASRProcessor", CohereASRProcessor) |