Add app.py and model

.gitignore

@@ -0,0 +1 @@
+.DS_Store

app.py

@@ -0,0 +1,91 @@
+import torch
+import jyutping
+from whisper_audio_classifier import WhisperAudioClassifier
+import librosa
+from transformers import WhisperFeatureExtractor
+
+feature_extractor = WhisperFeatureExtractor.from_pretrained(f"alvanlii/whisper-small-cantonese")
+feature_extractor.chunk_length = 3
+
+# Instantiate the model
+device = torch.device("mps")
+model = WhisperAudioClassifier().to(device)
+
+# Load the state dict
+state_dict = torch.load(f"whisper-small-encoder-bisyllabic-jyutping/checkpoints/model_epoch_1_step_1800.pth")
+
+# Load the state dict into the model
+model.load_state_dict(state_dict)
+
+# Set the model to evaluation mode
+model.eval()
+
+def predict(audio):
+    features = feature_extractor(audio, sampling_rate=16000)
+    with torch.no_grad():
+        inputs = torch.from_numpy(features['input_features'][0]).to(device)
+        inputs = inputs.unsqueeze(0)  # Add extra batch dimension in front
+        outs = model(inputs)
+        return [torch.softmax(tensor.squeeze(), dim=0).tolist() for tensor in outs]
+
+import gradio as gr
+import numpy as np
+
+def rank_initials(preds, k=3):
+    ranked = sorted([((jyutping.inflate_initial(i) if jyutping.inflate_initial(i) != '' else '∅'), p) for i, p in enumerate(preds)], key=lambda x: x[1], reverse=True)
+    return dict(ranked[:k])
+
+def rank_nucli(preds, k=3):
+    ranked = sorted([((jyutping.inflate_nucleus(i) if jyutping.inflate_nucleus(i) != '' else '∅'), p) for i, p in enumerate(preds)], key=lambda x: x[1], reverse=True)
+    return dict(ranked[:k])
+
+def rank_codas(preds, k=3):
+    ranked = sorted([((jyutping.inflate_coda(i) if jyutping.inflate_coda(i) != '' else '∅'), p) for i, p in enumerate(preds)], key=lambda x: x[1], reverse=True)
+    return dict(ranked[:k])
+
+def rank_tones(preds, k=3):
+    ranked = sorted([(str(i + 1), p) for i, p in enumerate(preds)], key=lambda x: x[1], reverse=True)
+    return dict(ranked[:k])
+
+def classify_audio(audio):
+    sampling_rate, audio = audio
+    audio = audio.astype(np.float32)
+    audio /= np.max(np.abs(audio))
+    audio_resampled = librosa.resample(audio, orig_sr=sampling_rate, target_sr=16000)
+    preds = predict(torch.from_numpy(audio_resampled))
+    return [
+        rank_initials(preds[0]),
+        rank_nucli(preds[1]),
+        rank_codas(preds[2]),
+        rank_tones(preds[3]),
+        rank_initials(preds[4]),
+        rank_nucli(preds[5]),
+        rank_codas(preds[6]),
+        rank_tones(preds[7]),
+    ]
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        inputs = gr.Audio(source="microphone", type="numpy", label="Input Audio")
+        submit_btn = gr.Button("Submit")
+
+    with gr.Row():
+        with gr.Column():
+            outputs_left = [
+                gr.Label(label="Initial 1"),
+                gr.Label(label="Nucleus 1"),
+                gr.Label(label="Coda 1"),
+                gr.Label(label="Tone 1"),
+            ]
+            
+        with gr.Column():
+            outputs_right = [
+                gr.Label(label="Initial 2"),
+                gr.Label(label="Nucleus 2"),
+                gr.Label(label="Coda 2"),
+                gr.Label(label="Tone 2"),
+            ]
+    
+    submit_btn.click(fn=classify_audio, inputs=inputs, outputs=outputs_left+outputs_right)
+
+demo.launch()

jyutping.py

@@ -0,0 +1,115 @@
+def extract_tone(syllable: str) -> int:
+    return int(syllable[-1]) - 1
+jyutping_initials = ['∅', 'ng', 'gw', 'kw', 'b', 'p', 'm', 'f', 'd', 't', 'n', 'l', 'g', 'k', 'h', 'w', 'z', 'c', 's', 'j']
+
+def extract_initial(syllable: str) -> (int, str):
+    for i, initial in enumerate(jyutping_initials):
+        if syllable.startswith(initial):
+            return (i, initial)
+    return (0, '')
+
+def inflate_initial(initial: int) -> str:
+    return jyutping_initials[initial] if initial != 0 else ''
+jyutping_nuclei = ['∅', 'aa', 'yu', 'eo', 'oe', 'a', 'i', 'u', 'e', 'o']
+
+def extract_nucleus(syllable: str, initial: str) -> int:
+    syllable = syllable[len(initial):]
+    for i, nucleus in enumerate(jyutping_nuclei):
+        if syllable.startswith(nucleus):
+            return (i, nucleus)
+    return (0, '')
+
+def inflate_nucleus(nucleus: int) -> str:
+    return jyutping_nuclei[nucleus] if nucleus != 0 else ''
+jyutping_codas = ['∅', 'ng', 'p', 't', 'k', 'm', 'n', 'i', 'u']
+
+def extract_coda(syllable: str, initial: str, nucleus: str) -> int:
+    syllable = syllable[len(initial) + len(nucleus):]
+    for i, coda in enumerate(jyutping_codas):
+        if syllable.startswith(coda):
+            return (i, coda)
+    return (0, '')
+
+def inflate_coda(coda: int) -> str:
+    return jyutping_codas[coda] if coda != 0 else ''
+syllable = 'neoi5'
+
+initial_int, initial = extract_initial(syllable)
+nucleus_int, nucleus = extract_nucleus(syllable, initial)
+coda_int, coda = extract_coda(syllable, initial, nucleus)
+
+assert(initial == 'n' and initial_int == 10)
+assert(nucleus == 'eo' and nucleus_int == 3)
+assert(coda == 'i' and coda_int == 7)
+syllable = 'gwok3'
+
+initial_int, initial = extract_initial(syllable)
+nucleus_int, nucleus = extract_nucleus(syllable, initial)
+coda_int, coda = extract_coda(syllable, initial, nucleus)
+
+assert(initial == 'gw' and initial_int == 2)
+assert(nucleus == 'o' and nucleus_int == 9)
+assert(coda == 'k' and coda_int == 4)
+syllable = 'oi3'
+
+initial_int, initial = extract_initial(syllable)
+nucleus_int, nucleus = extract_nucleus(syllable, initial)
+coda_int, coda = extract_coda(syllable, initial, nucleus)
+
+assert(initial == '' and initial_int == 0)
+assert(nucleus == 'o' and nucleus_int == 9)
+assert(coda == 'i' and coda_int == 7)
+syllable = 'ng4'
+
+initial_int, initial = extract_initial(syllable)
+nucleus_int, nucleus = extract_nucleus(syllable, initial)
+coda_int, coda = extract_coda(syllable, initial, nucleus)
+
+assert(initial == 'ng' and initial_int == 1)
+assert(nucleus == '' and nucleus_int == 0)
+assert(coda == '' and coda_int == 0)
+syllable = 'ngo5'
+
+initial_int, initial = extract_initial(syllable)
+nucleus_int, nucleus = extract_nucleus(syllable, initial)
+coda_int, coda = extract_coda(syllable, initial, nucleus)
+
+assert(initial == 'ng' and initial_int == 1)
+assert(nucleus == 'o' and nucleus_int == 9)
+assert(coda == '' and coda_int == 0)
+syllable = 'a3'
+
+initial_int, initial = extract_initial(syllable)
+nucleus_int, nucleus = extract_nucleus(syllable, initial)
+coda_int, coda = extract_coda(syllable, initial, nucleus)
+
+assert(initial == '' and initial_int == 0)
+assert(nucleus == 'a' and nucleus_int == 5)
+assert(coda == '' and coda_int == 0)
+syllable = 'aa3'
+
+initial_int, initial = extract_initial(syllable)
+nucleus_int, nucleus = extract_nucleus(syllable, initial)
+coda_int, coda = extract_coda(syllable, initial, nucleus)
+
+assert(initial == '' and initial_int == 0)
+assert(nucleus == 'aa' and nucleus_int == 1)
+assert(coda == '' and coda_int == 0)
+syllable = 'ngaang6'
+
+initial_int, initial = extract_initial(syllable)
+nucleus_int, nucleus = extract_nucleus(syllable, initial)
+coda_int, coda = extract_coda(syllable, initial, nucleus)
+
+assert(initial == 'ng' and initial_int == 1)
+assert(nucleus == 'aa' and nucleus_int == 1)
+assert(coda == 'ng' and coda_int == 1)
+def extract_jyutping(syllable: str) -> (int, int, int, int):
+    initial_int, initial = extract_initial(syllable)
+    nucleus_int, nucleus = extract_nucleus(syllable, initial)
+    coda_int, _ = extract_coda(syllable, initial, nucleus)
+    tone = extract_tone(syllable)
+    return (initial_int, nucleus_int, coda_int, tone)
+
+def inflate_jyutping(initial: int, nucleus: int, coda: int, tone: int) -> str:
+    return f"{inflate_initial(initial)}{inflate_nucleus(nucleus)}{inflate_coda(coda)}{tone + 1}"

requirements.txt

@@ -0,0 +1,3 @@
+transformers
+torch
+librosa

whisper-small-encoder-bisyllabic-jyutping/checkpoints/model_epoch_1_step_1800.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5f85b07f5287b93c473b234696b387a6b8ff1414bc1980b811f7933f9ecddb28
+size 390773292

whisper_audio_classifier.py

@@ -0,0 +1,69 @@
+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