hubert phoneme + quick test model

app/app.py

@@ -4,7 +4,7 @@
 import gradio as gr
 import pandas as pd
 
-from tasks import start_eval_task, get_status
+from tasks import start_eval_task, get_status, run_sample_inference
 from hf import get_or_create_leaderboard
 
 from codes import CODES
@@ -205,6 +205,16 @@ with gr.Blocks(
                 outputs=result,
             )
 
+            gr.Markdown("---\n### Test Model")
+            test_audio = gr.Audio(interactive=True, format="wav")
+            test_btn = gr.Button("Run")
+            test_result = gr.Textbox(label="Test Result")
+            test_btn.click(
+                fn=run_sample_inference,
+                inputs=[test_audio, model_id, model_type, output_code],
+                outputs=test_result,
+            )
+
         with gr.TabItem("📊 Submission Status"):
             query = gr.Textbox(
                 label="Model ID or Task ID",

app/inference.py

@@ -3,8 +3,9 @@
 import torch
 from transformers import AutoProcessor, AutoModelForCTC
 from espnet2.bin.s2t_inference import Speech2Text
+from inference_huberphoneme import HuBERTPhoneme, Tokenizer
 
-MODEL_TYPES = ["Transformers CTC", "POWSM"]
+MODEL_TYPES = ["Transformers CTC", "POWSM", "HuBERTPhoneme"]
 
 DEVICE = (
     "cuda"
@@ -78,6 +79,23 @@ def transcribe_transformers_ctc(audio, model) -> str:
     return processor.decode(predicted_ids[0])
 
 
+# ===========================================================================
+# ============================== HuBERTPhoneme ==============================
+def load_hubert_phoneme(model_id, device=DEVICE):
+    model = HuBERTPhoneme.from_pretrained(model_id).to(device).eval()
+    tokenizer = Tokenizer(with_blank=model.ctc_training)
+    return model, tokenizer, device
+
+
+def transcribe_hubert_phoneme(audio, model) -> str:
+    model, tokenizer, device = model
+    with torch.inference_mode():
+        output, _ = model.inference(torch.from_numpy(audio).to(device).unsqueeze(0))
+        predictions = output.argmax(dim=-1).squeeze().cpu()
+        arpabet = tokenizer.decode(predictions.unique_consecutive())
+        return arpabet
+
+
 # ===========================================================================
 
 
@@ -86,6 +104,8 @@ def load_model(model_id, type, device=DEVICE):
         return load_powsm(model_id, device=device)
     elif type == "Transformers CTC":
         return load_transformers_ctc(model_id, device=device)
+    elif type == "HuBERTPhoneme":
+        return load_hubert_phoneme(model_id, device=device)
     else:
         raise ValueError("Unsupported model type: " + str(type))
 
@@ -95,5 +115,7 @@ def transcribe(audio, type, model) -> str:
         return transcribe_powsm(audio, model)
     elif type == "Transformers CTC":
         return transcribe_transformers_ctc(audio, model)
+    elif type == "HuBERTPhoneme":
+        return transcribe_hubert_phoneme(audio, model)
     else:
         raise ValueError("Unsupported model type: " + str(type))

app/inference_huberphoneme.py

@@ -0,0 +1,133 @@
+# Adapted from https://github.com/bootphon/spokenlm-phoneme
+
+import torch
+import torchaudio
+from huggingface_hub import PyTorchModelHubMixin
+from torch import Tensor, nn
+from torchaudio.models.wav2vec2 import components
+from torchaudio.pipelines import HUBERT_BASE
+from typing import Iterable
+
+
+class Tokenizer:
+    # fmt:off
+    PHONEMES = {
+        "SIL": 0, "AA": 1, "AE": 2, "AH": 3, "AO": 4, "AW": 5, "AY": 6, "B": 7,
+        "CH": 8, "D": 9, "DH": 10, "EH": 11, "ER": 12, "EY": 13, "F": 14, "G": 15,
+        "HH": 16, "IH": 17, "IY": 18, "JH": 19, "K": 20, "L": 21, "M": 22, "N": 23,
+        "NG": 24, "OW": 25, "OY": 26, "P": 27, "R": 28, "S": 29, "SH": 30, "T": 31,
+        "TH": 32, "UH": 33, "UW": 34, "V": 35, "W": 36, "Y": 37, "Z": 38, "ZH": 39,
+    }
+    # fmt:on
+
+    def __init__(self, with_blank: bool = False) -> None:
+        self.token_to_id = self.PHONEMES | {"<pad>": self.pad_id}
+        self.id_to_token = {v: k for k, v in self.token_to_id.items()}
+        self.with_blank = with_blank
+
+    @property
+    def vocab_size(self) -> int:
+        if self.with_blank:
+            return len(self.PHONEMES) + 1
+        return len(self.PHONEMES)
+
+    @property
+    def silence_id(self) -> int:
+        return self.PHONEMES["SIL"]
+
+    @property
+    def pad_id(self) -> int:
+        return len(self.PHONEMES)
+
+    def encode(self, phones: "list[str] | str") -> torch.LongTensor:
+        if isinstance(phones, str):
+            phones = phones.split(" ")
+        return torch.LongTensor([self.token_to_id[phone] for phone in phones])
+
+    def decode(self, tokens: Iterable[int]) -> str:
+        return " ".join(
+            self.id_to_token[int(token)]
+            for token in tokens
+            if token < self.pad_id and int(token) != self.silence_id
+        )
+
+
+FINETUNING_HUBERT_CONFIG = {
+    "encoder_projection_dropout": 0,
+    "encoder_attention_dropout": 0,
+    "encoder_ff_interm_dropout": 0.1,
+    "encoder_dropout": 0,
+    "encoder_layer_drop": 0.1,  # In torchaudio: 0.05
+    "mask_prob": 0.75,  # In torchaudio: 0.65
+    "mask_channel_prob": 0.5,
+    "mask_channel_length": 10,  # In torchaudio and fairseq: 64. This is the value for pretraining.
+    "num_classes": 500,  # Number of classes during HuBERT pretraining.
+}
+
+
+class HuBERTPhoneme(nn.Module, PyTorchModelHubMixin):
+    def __init__(self, freeze_encoder: bool = True, ctc_training: bool = False) -> None:
+        """Initialize the model.
+
+        Parameters
+        ----------
+        freeze_encoder : bool, optional
+            Whether to freeze the Transformer encoder of HuBERT, by default True.
+            The convolutional layers are always frozen.
+        """
+        super().__init__()
+        self.model = torchaudio.models.hubert_pretrain_base(**FINETUNING_HUBERT_CONFIG)
+        self.model.wav2vec2.load_state_dict(HUBERT_BASE.get_model().state_dict())
+        self.aux = nn.Linear(
+            HUBERT_BASE._params["encoder_embed_dim"],
+            Tokenizer(with_blank=ctc_training).vocab_size,
+        )
+        self.freeze_encoder = freeze_encoder
+        self.ctc_training = ctc_training
+
+    def forward(
+        self, waveforms: Tensor, lengths: "Tensor | None" = None
+    ) -> "tuple[Tensor, Tensor | None]":
+        """Extract logits during training, with masking."""
+        if self.freeze_encoder:
+            with torch.no_grad():
+                x, out_len = self.model.wav2vec2.feature_extractor(waveforms, lengths)
+                padding_mask = components._get_padding_mask(x, out_len)
+                x, attention_mask = self.model.wav2vec2.encoder._preprocess(x, out_len)  # type: ignore
+                x, _ = self.model.mask_generator(x, padding_mask)
+                x = self.model.wav2vec2.encoder.transformer(x, attention_mask=attention_mask)  # type: ignore
+        else:
+            with torch.no_grad():
+                x, out_len = self.model.wav2vec2.feature_extractor(waveforms, lengths)
+                padding_mask = components._get_padding_mask(x, out_len)
+            x, attention_mask = self.model.wav2vec2.encoder._preprocess(x, out_len)  # type: ignore
+            x, _ = self.model.mask_generator(x, padding_mask)
+            x = self.model.wav2vec2.encoder.transformer(x, attention_mask=attention_mask)  # type: ignore
+        logits = self.aux(x)
+        return logits, out_len
+
+    def inference(
+        self, waveforms: Tensor, lengths: "Tensor | None" = None
+    ) -> "tuple[Tensor, Tensor | None]":
+        """Extract logits during inference. No masking is applied."""
+        x, out_len = self.model.wav2vec2(waveforms, lengths)
+        logits = self.aux(x)
+        return logits, out_len
+
+    @torch.jit.export
+    def extract_features(
+        self, waveforms: Tensor, lengths: "Tensor | None" = None
+    ) -> "tuple[list[Tensor], Tensor | None]":
+        """Extract features from intermediate layers. No masking is applied."""
+        x, out_len = self.model.wav2vec2.extract_features(waveforms, lengths)
+        x.append(self.aux(x[-1]))
+        return x, out_len
+
+    def train(self, mode: bool = True) -> "HuBERTPhoneme":
+        """Override the train method to set the encoder in eval mode if it is frozen."""
+        if self.freeze_encoder:
+            self.model.wav2vec2.eval()
+        else:
+            self.model.wav2vec2.train(mode)
+        self.aux.train(mode)
+        return self

app/tasks.py

@@ -5,6 +5,8 @@ import multiprocessing
 from typing import TypedDict
 from datetime import datetime
 
+import librosa
+import numpy as np
 
 from metrics import per, fer
 from datasets import load_from_disk
@@ -127,3 +129,25 @@ def _eval_task(task: Task, leaderboard_lock):
     except Exception as e:
         task["status"] = "failed"
         task["error"] = str(e)
+
+
+def run_sample_inference(audio, model_id: str, model_type: str, phone_code: str):
+    clear_cache()
+
+    # Load model
+    model = load_model(model_id, model_type)
+
+    # Format audio as monochannel 16 kHz float32
+    sample_rate, wav_array = audio
+    wav_array = wav_array.astype(np.float32)
+    if wav_array.ndim == 2 and wav_array.shape[1] == 2:
+        wav_array = np.mean(wav_array, axis=1)
+    wav_array = librosa.resample(y=wav_array, orig_sr=sample_rate, target_sr=16_000)
+
+    # Transcribe
+    transcript = transcribe(wav_array, model_type, model)
+    if phone_code != "ipa":
+        transcript = convert(transcript, phone_code, "ipa")
+
+    clear_cache()
+    return transcript

requirements.txt

@@ -6,9 +6,9 @@ datasets==4.0.0
 pandas==2.3.3
 numpy==2.0.2
 panphon==0.21.2
-torch==2.8.0
-torchaudio==2.8.0
-torchcodec==0.6.0
+torch==2.9.1
+torchaudio==2.9.1
+torchcodec==0.8.0
 transformers==4.56.0
 phonemizer==3.3.0
 espnet==202509
@@ -17,3 +17,4 @@ espnet-model-zoo==0.1.7
 # UI
 gradio==5.12.0
 protobuf==6.32.0
+pydantic==2.10.6

requirements_lock.txt

@@ -90,8 +90,8 @@ propcache==0.3.2
 protobuf==6.32.0
 pyarrow==21.0.0
 pycparser==2.23
-pydantic==2.11.7
-pydantic_core==2.33.2
+pydantic==2.10.6
+pydantic_core==2.27.2
 pydub==0.25.1
 Pygments==2.19.2
 pyparsing==3.2.3
@@ -127,10 +127,10 @@ sympy==1.14.0
 threadpoolctl==3.6.0
 tokenizers==0.22.0
 tomlkit==0.13.3
-torch==2.8.0
+torch==2.9.1
 torch-complex==0.4.4
-torchaudio==2.8.0
-torchcodec==0.6.0
+torchaudio==2.9.1
+torchcodec==0.8.0
 torchmetrics==1.8.2
 tqdm==4.67.1
 transformers==4.56.0