| """Batched GPU log-mel matching transformers' Gemma4AudioFeatureExtractor, plus SpecAugment.""" |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from transformers import AutoFeatureExtractor |
|
|
|
|
| class LogMel(nn.Module): |
| def __init__(self, model_id: str): |
| super().__init__() |
| fe = AutoFeatureExtractor.from_pretrained(model_id) |
| assert fe.preemphasis == 0.0 and fe.dither == 0.0 and fe.input_scale_factor == 1.0 |
| assert fe.per_bin_mean is None and fe.per_bin_stddev is None |
| self.frame_length, self.hop_length, self.fft_length = fe.frame_length, fe.hop_length, fe.fft_length |
| self.mel_floor = float(fe.mel_floor) |
| self.register_buffer("window", torch.tensor(fe.window, dtype=torch.float32), persistent=False) |
| self.register_buffer("mel_filters", torch.tensor(fe.mel_filters, dtype=torch.float32), persistent=False) |
|
|
| @torch.no_grad() |
| def forward(self, wav: torch.Tensor, wav_lens: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| """wav (B, L) float32 zero-padded -> feats (B, T, 128), mask (B, T).""" |
| pad_left = self.frame_length // 2 |
| wav = F.pad(wav.float(), (pad_left, 0)) |
| frames = wav.unfold(-1, self.frame_length + 1, self.hop_length)[..., :-1] |
| spec = torch.fft.rfft(frames * self.window, n=self.fft_length, dim=-1).abs() |
| feats = torch.log(spec @ self.mel_filters + self.mel_floor) |
| |
| frame_end = torch.arange(feats.shape[1], device=wav.device) * self.hop_length + self.frame_length |
| mask = frame_end[None, :] < (wav_lens[:, None] + pad_left) |
| return feats * mask[..., None], mask |
|
|
|
|
| @torch.no_grad() |
| def spec_augment(feats, mask, freq_masks=2, freq_width=27, time_masks=10, time_width_frac=0.04): |
| """In-place-free SpecAugment on (B, T, F) log-mel; masked bins are set to the utterance mean.""" |
| B, T, Fdim = feats.shape |
| dev = feats.device |
| lens = mask.sum(1) |
| fill = (feats * mask[..., None]).sum((1, 2)) / (lens * Fdim).clamp(min=1) |
| keep = torch.ones(B, T, Fdim, dtype=torch.bool, device=dev) |
| f_idx = torch.arange(Fdim, device=dev) |
| t_idx = torch.arange(T, device=dev) |
| for _ in range(freq_masks): |
| w = torch.randint(0, freq_width + 1, (B,), device=dev) |
| s = (torch.rand(B, device=dev) * (Fdim - w)).long() |
| keep &= ~((f_idx[None] >= s[:, None]) & (f_idx[None] < (s + w)[:, None]))[:, None, :] |
| max_w = (lens.float() * time_width_frac).long() |
| for _ in range(time_masks): |
| w = (torch.rand(B, device=dev) * (max_w + 1)).long() |
| s = (torch.rand(B, device=dev) * (lens - w).clamp(min=1)).long() |
| keep &= ~((t_idx[None] >= s[:, None]) & (t_idx[None] < (s + w)[:, None]))[:, :, None] |
| out = torch.where(keep, feats, fill[:, None, None]) |
| return out * mask[..., None] |
|
|
