Spaces:
Sleeping
Sleeping
File size: 3,825 Bytes
d4c10cd |
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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 |
import os
import time
import torch
import torchaudio
import gradio as gr
from transformers import (
Wav2Vec2Processor, HubertForCTC,
WhisperProcessor, WhisperForConditionalGeneration
)
from phonemizer import phonemize
import difflib
# === Setup: Load all 3 models ===
# 1. Base HuBERT
base_proc = Wav2Vec2Processor.from_pretrained("facebook/hubert-large-ls960-ft")
base_model = HubertForCTC.from_pretrained("facebook/hubert-large-ls960-ft").eval()
# 2. Whisper + phonemizer
whisper_proc = WhisperProcessor.from_pretrained("openai/whisper-base")
whisper_model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-base").eval()
# 3. My Hubert Model
token = os.environ.get("HF_TOKEN")
your_proc = Wav2Vec2Processor.from_pretrained("tecasoftai/hubert-finetune", token=token)
your_model = HubertForCTC.from_pretrained("tecasoftai/hubert-finetune", token=token).eval()
# === Helper ===
def load_audio(filepath):
waveform, sr = torchaudio.load(filepath)
if sr != 16000:
waveform = torchaudio.functional.resample(waveform, sr, 16000)
return waveform.squeeze()
def calc_per(pred, ref):
pred_list = pred.strip().split()
ref_list = ref.strip().split()
sm = difflib.SequenceMatcher(None, ref_list, pred_list)
dist = sum(tr[-1] for tr in sm.get_opcodes() if tr[0] != 'equal')
if len(ref_list) == 0:
return 0.0
return round(100 * dist / len(ref_list), 2)
# === Inference functions ===
def run_hubert_base(wav):
start = time.time()
inputs = base_proc(wav, sampling_rate=16000, return_tensors="pt")
with torch.no_grad():
logits = base_model(**inputs).logits
ids = torch.argmax(logits, dim=-1)
phonemes = base_proc.batch_decode(ids)[0]
return phonemes, time.time() - start
def run_whisper(wav):
start = time.time()
inputs = whisper_proc(wav, sampling_rate=16000, return_tensors="pt")
with torch.no_grad():
ids = whisper_model.generate(inputs["input_features"])
text = whisper_proc.batch_decode(ids, skip_special_tokens=True)[0]
phonemes = phonemize(text, language='en-us', backend='espeak')
return phonemes, time.time() - start
def run_your_model(wav):
start = time.time()
inputs = your_proc(wav, sampling_rate=16000, return_tensors="pt")
with torch.no_grad():
logits = your_model(**inputs).logits
ids = torch.argmax(logits, dim=-1)
phonemes = your_proc.batch_decode(ids)[0]
return phonemes, time.time() - start
# === Main Gradio function ===
def benchmark_all(audio_path, reference_phoneme):
wav = load_audio(audio_path)
results = []
# 1. HuBERT Base
phonemes, dur = run_hubert_base(wav)
per = calc_per(phonemes, reference_phoneme)
results.append(["HuBERT-Base", phonemes, f"{dur:.2f}s", f"{per}%"])
# 2. Whisper
phonemes, dur = run_whisper(wav)
per = calc_per(phonemes, reference_phoneme)
results.append(["Whisper + Phonemizer", phonemes, f"{dur:.2f}s", f"{per}%"])
# 3. My Hubert model
phonemes, dur = run_your_model(wav)
per = calc_per(phonemes, reference_phoneme)
results.append(["Your HuBERT (fine-tuned)", phonemes, f"{dur:.2f}s", f"{per}%"])
return results
# === UI ===
demo = gr.Interface(
fn=benchmark_all,
inputs=[
gr.Audio(type="filepath", label="Upload Audio"),
gr.Textbox(label="Ground-truth Phonemes (space-separated)", placeholder="f ə n ə m aɪ z")
],
outputs=gr.Dataframe(headers=["Model", "Phoneme Output", "Inference Time", "PER (%)"]),
title="Phoneme Recognition Benchmark",
description="Compare HuBERT-Base, Whisper, and your fine-tuned model on phoneme recognition."
)
if __name__ == "__main__":
demo.launch()
|