Spaces:
Sleeping
Sleeping
| """Shared constants and the audio feature frontend for tiny-kws. | |
| Everything that must be identical between training, evaluation and the live | |
| demo lives here: the label order, the feature parameters, and the log-mel | |
| transform itself. If train and inference disagree on any of these, accuracy | |
| silently collapses, so they are defined exactly once. | |
| """ | |
| import json | |
| from pathlib import Path | |
| import torch | |
| import torch.nn as nn | |
| import torchaudio | |
| # --------------------------------------------------------------------------- | |
| # Task definition: the standard 12-class Speech Commands benchmark | |
| # --------------------------------------------------------------------------- | |
| KEYWORDS = ["yes", "no", "up", "down", "left", "right", "on", "off", "stop", "go"] | |
| LABELS = ["silence", "unknown"] + KEYWORDS # index 0 = silence, 1 = unknown | |
| LABEL_TO_IDX = {l: i for i, l in enumerate(LABELS)} | |
| # The 25 non-keyword words in Speech Commands v2; clips from these become "unknown". | |
| UNKNOWN_WORDS = [ | |
| "backward", "bed", "bird", "cat", "dog", "eight", "five", "follow", | |
| "forward", "four", "happy", "house", "learn", "marvin", "nine", "one", | |
| "seven", "sheila", "six", "three", "tree", "two", | |
| "visual", "wow", "zero", | |
| ] | |
| # --------------------------------------------------------------------------- | |
| # Audio / feature parameters | |
| # --------------------------------------------------------------------------- | |
| SAMPLE_RATE = 16_000 # Hz; dataset native rate | |
| CLIP_SAMPLES = 16_000 # exactly 1 second per clip | |
| N_FFT = 400 # 25 ms analysis window | |
| HOP_LENGTH = 160 # 10 ms hop -> 101 frames per 1 s clip | |
| N_MELS = 64 | |
| F_MIN, F_MAX = 20.0, 7600.0 # stay below Nyquist (8 kHz) | |
| N_FRAMES = CLIP_SAMPLES // HOP_LENGTH + 1 # 101 (torch STFT uses center padding) | |
| LOG_EPS = 1e-6 | |
| class LogMel(nn.Module): | |
| """Waveform (B, 16000) -> log-mel spectrogram (B, 1, 64, 101). | |
| Kept as an nn.Module so it can be moved to CUDA for fast feature | |
| extraction on Colab; on Apple Silicon it should stay on CPU because | |
| torchaudio's spectrogram ops have gaps on MPS. | |
| """ | |
| def __init__(self): | |
| super().__init__() | |
| self.mel = torchaudio.transforms.MelSpectrogram( | |
| sample_rate=SAMPLE_RATE, | |
| n_fft=N_FFT, | |
| hop_length=HOP_LENGTH, | |
| n_mels=N_MELS, | |
| f_min=F_MIN, | |
| f_max=F_MAX, | |
| power=2.0, | |
| ) | |
| def forward(self, wav: torch.Tensor) -> torch.Tensor: | |
| if wav.dim() == 1: | |
| wav = wav.unsqueeze(0) | |
| feats = torch.log(self.mel(wav) + LOG_EPS) | |
| return feats.unsqueeze(1) # add channel dim | |
| def normalize(feats: torch.Tensor, stats: dict) -> torch.Tensor: | |
| """Apply train-set global mean/std normalization.""" | |
| return (feats - stats["mean"]) / stats["std"] | |
| def load_stats(path: Path) -> dict: | |
| with open(path) as f: | |
| return json.load(f) | |
| def pick_device() -> torch.device: | |
| if torch.cuda.is_available(): | |
| return torch.device("cuda") | |
| if torch.backends.mps.is_available(): | |
| return torch.device("mps") | |
| return torch.device("cpu") | |