|
|
import whisper |
|
|
from torch import Tensor |
|
|
import torch.nn.functional as F |
|
|
from whisper.model import AudioEncoder |
|
|
|
|
|
|
|
|
class WhisperAudioEncoder(AudioEncoder): |
|
|
""" |
|
|
We inherited the original Whisper encoder and modified its 30-second fixed-length padding logic to |
|
|
improve training and inference efficiency. |
|
|
""" |
|
|
def __init__(self, n_mels: int, n_ctx: int, n_state: int, n_head: int, n_layer: int): |
|
|
super().__init__(n_mels, n_ctx, n_state, n_head, n_layer) |
|
|
self.audio_emb_dim = n_state |
|
|
|
|
|
def forward(self, x: Tensor): |
|
|
""" |
|
|
x : torch.Tensor, shape = [B, T, d] the mel spectrogram of the audio |
|
|
""" |
|
|
x = x.transpose(1, 2) |
|
|
x = F.gelu(self.conv1(x)) |
|
|
x = F.gelu(self.conv2(x)) |
|
|
x = x.permute(0, 2, 1) |
|
|
positional_embedding = self.positional_embedding[:x.shape[1], :] |
|
|
assert x.shape[1:] == positional_embedding.shape, "incorrect audio shape" |
|
|
x = (x + positional_embedding).to(x.dtype) |
|
|
|
|
|
for block in self.blocks: |
|
|
x = block(x) |
|
|
|
|
|
x = self.ln_post(x) |
|
|
return x |
|
|
|
|
|
@classmethod |
|
|
def from_pretrained(cls, model_path, **kwargs): |
|
|
whisper_model = whisper.load_model(model_path) |
|
|
audio_encoder = cls( |
|
|
whisper_model.dims.n_mels, |
|
|
whisper_model.dims.n_audio_ctx*10, |
|
|
whisper_model.dims.n_audio_state, |
|
|
whisper_model.dims.n_audio_head, |
|
|
whisper_model.dims.n_audio_layer, |
|
|
).to(whisper_model.device) |
|
|
state_dict = whisper_model.encoder.state_dict() |
|
|
state_dict.pop('positional_embedding') |
|
|
ret = audio_encoder.load_state_dict(state_dict, strict=False) |
|
|
logger.warning(f'whisper encoder does not load `positional_embedding`. {ret}') |
|
|
audio_encoder.audio_emb_dim = whisper_model.dims.n_audio_state |
|
|
return audio_encoder |
