Spaces:
Running
Running
| from transformers import WhisperModel | |
| from torch import nn | |
| import torch | |
| from jyutping import jyutping_initials, jyutping_nuclei, jyutping_codas | |
| class WhisperAudioClassifier(nn.Module): | |
| def __init__(self): | |
| super(WhisperAudioClassifier, self).__init__() | |
| # Load the Whisper model encoder | |
| self.whisper_encoder = WhisperModel.from_pretrained(f"alvanlii/whisper-small-cantonese", device_map="auto").get_encoder() | |
| self.whisper_encoder.eval() # Set the Whisper model to evaluation mode | |
| # Assuming we know the output size of the Whisper encoder, or it needs to be determined | |
| whisper_output_size = 768 | |
| self.tone_attention = nn.MultiheadAttention(whisper_output_size, 8, dropout=0.1, batch_first=True) | |
| self.initial_attention = nn.MultiheadAttention(whisper_output_size, 8, dropout=0.1, batch_first=True) | |
| self.nucleus_attention = nn.MultiheadAttention(whisper_output_size, 8, dropout=0.1, batch_first=True) | |
| self.coda_attention = nn.MultiheadAttention(whisper_output_size, 8, dropout=0.1, batch_first=True) | |
| self.pool = nn.AdaptiveAvgPool1d(1) | |
| # Separate output layers for each class set | |
| self.initial_fc1 = nn.Linear(whisper_output_size, len(jyutping_initials)) | |
| self.nucleus_fc1 = nn.Linear(whisper_output_size, len(jyutping_nuclei)) | |
| self.coda_fc1 = nn.Linear(whisper_output_size, len(jyutping_codas)) | |
| self.tone_fc1 = nn.Linear(whisper_output_size, 6) | |
| self.initial_fc2 = nn.Linear(whisper_output_size, len(jyutping_initials)) | |
| self.nucleus_fc2 = nn.Linear(whisper_output_size, len(jyutping_nuclei)) | |
| self.coda_fc2 = nn.Linear(whisper_output_size, len(jyutping_codas)) | |
| self.tone_fc2 = nn.Linear(whisper_output_size, 6) | |
| self.dropout = nn.Dropout(0.1) | |
| def forward(self, x): | |
| # Use Whisper model to encode audio input | |
| with torch.no_grad(): # No need to track gradients for the encoder | |
| x = self.whisper_encoder(x).last_hidden_state | |
| initial, _ = self.initial_attention(x, x, x, need_weights=False) | |
| initial = initial.permute(0, 2, 1) # [batch_size, channels, seq_len] | |
| initial = self.pool(initial) # [batch_size, channels, 1] | |
| initial = initial.squeeze(-1) # [batch_size, channels] | |
| initial_out1 = self.initial_fc1(initial) | |
| initial_out2 = self.initial_fc2(initial) | |
| nucleus, _ = self.nucleus_attention(x, x, x, need_weights=False) | |
| nucleus = nucleus.permute(0, 2, 1) # [batch_size, channels, seq_len] | |
| nucleus = self.pool(nucleus) # [batch_size, channels, 1] | |
| nucleus = nucleus.squeeze(-1) # [batch_size, channels] | |
| nucleus_out1 = self.nucleus_fc1(nucleus) | |
| nucleus_out2 = self.nucleus_fc2(nucleus) | |
| coda, _ = self.coda_attention(x, x, x, need_weights=False) | |
| coda = coda.permute(0, 2, 1) # [batch_size, channels, seq_len] | |
| coda = self.pool(coda) # [batch_size, channels, 1] | |
| coda = coda.squeeze(-1) # [batch_size, channels] | |
| coda_out1 = self.coda_fc1(coda) | |
| coda_out2 = self.coda_fc2(coda) | |
| tone, _ = self.tone_attention(x, x, x, need_weights=False) | |
| tone = tone.permute(0, 2, 1) # [batch_size, channels, seq_len] | |
| tone = self.pool(tone) # [batch_size, channels, 1] | |
| tone = tone.squeeze(-1) # [batch_size, channels] | |
| tone_out1 = self.tone_fc1(tone) | |
| tone_out2 = self.tone_fc2(tone) | |
| return initial_out1, nucleus_out1, coda_out1, tone_out1, initial_out2, nucleus_out2, coda_out2, tone_out2 | |