File size: 3,438 Bytes
14fbf03 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | """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 # 128000 samples -> 800 mel frames -> 400 positions
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) # [80, 201]
def forward(self, x): # x [B, N_SAMPLES] float
st = torch.stft(x.float(), NFFT, HOP, window=self.window, return_complex=True)
mag = st[..., :-1].abs() ** 2 # [B, 201, 800]
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 # [B, 400, d]
a = self.pool(h).softmax(1)
return self.head((a * h).sum(1)).squeeze(-1) # logit
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() # [402, 1, 400]
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) # [B, 402, 801]
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
|