| """End-of-turn classifier: Whisper encoder (base for the released model) on the last 8 s of audio -> P(speaker is done).""" |
| import math, torch, torch.nn as nn |
| from transformers import WhisperConfig, WhisperModel |
| from transformers.models.whisper.modeling_whisper import WhisperEncoder |
| SR, CTX_S, N_MELS, HOP, NFFT = 16000, 8, 80, 160, 400 |
| N_SAMPLES = SR * CTX_S; N_FRAMES = N_SAMPLES // HOP |
| class LogMel(nn.Module): |
| """Whisper's log-mel front end in torch (matches WhisperFeatureExtractor up to float error).""" |
| def __init__(self): |
| super().__init__() |
| from transformers import WhisperFeatureExtractor |
| fe = WhisperFeatureExtractor(feature_size=N_MELS) |
| self.register_buffer("window", torch.hann_window(NFFT), persistent=False) |
| self.register_buffer("fb", torch.tensor(fe.mel_filters, dtype=torch.float32).T, persistent=False) |
| def forward(self, x): |
| st = torch.stft(x.float(), NFFT, HOP, window=self.window, return_complex=True) |
| mag = st[..., :-1].abs() ** 2 |
| mel = torch.clamp(self.fb @ mag, min=1e-10).log10() |
| mel = torch.maximum(mel, mel.amax(dim=(1, 2), keepdim=True) - 8.0) |
| return (mel + 4.0) / 4.0 |
| class EOT(nn.Module): |
| def __init__(self, base="openai/whisper-tiny", pretrained=True): |
| super().__init__() |
| cfg = WhisperConfig.from_pretrained(base); cfg.max_source_positions = N_FRAMES // 2 |
| self.mel = LogMel(); self.enc = WhisperEncoder(cfg) |
| if pretrained: |
| sd = WhisperModel.from_pretrained(base).encoder.state_dict() |
| sd["embed_positions.weight"] = sd["embed_positions.weight"][: cfg.max_source_positions] |
| self.enc.load_state_dict(sd) |
| d = cfg.d_model |
| self.pool = nn.Linear(d, 1) |
| self.head = nn.Sequential(nn.Linear(d, 256), nn.GELU(), nn.Dropout(0.1), nn.Linear(256, 1)) |
| def forward_mel(self, mel): |
| h = self.enc(mel).last_hidden_state |
| a = self.pool(h).softmax(1) |
| return self.head((a * h).sum(1)).squeeze(-1) |
| def forward(self, wav): |
| return self.forward_mel(self.mel(wav)) |
|
|
|
|
| class ConvLogMel(LogMel): |
| """Same as LogMel but the STFT is a fixed DFT conv1d (center=True, reflect pad): exports to any ONNX runtime.""" |
| def __init__(self): |
| super().__init__() |
| n = torch.arange(NFFT, dtype=torch.float64); k = torch.arange(NFFT // 2 + 1, dtype=torch.float64)[:, None] |
| ang = 2 * math.pi * k * n / NFFT; win = torch.hann_window(NFFT, dtype=torch.float64) |
| basis = torch.cat([torch.cos(ang) * win, -torch.sin(ang) * win])[:, None, :].float() |
| self.register_buffer("basis", basis, persistent=False) |
| def forward(self, x): |
| x = torch.nn.functional.pad(x.float()[:, None, :], (NFFT // 2, NFFT // 2), mode="reflect") |
| st = torch.nn.functional.conv1d(x, self.basis, stride=HOP) |
| re, im = st[:, : NFFT // 2 + 1], st[:, NFFT // 2 + 1:] |
| mag = (re ** 2 + im ** 2)[..., :-1] |
| mel = torch.clamp(self.fb @ mag, min=1e-10).log10() |
| mel = torch.maximum(mel, mel.amax(dim=(1, 2), keepdim=True) - 8.0) |
| return (mel + 4.0) / 4.0 |
|
|
