#!/usr/bin/env python3 """ Pronunciation Trainer — HF Spaces A random English word is shown with its reference IPA. The user records speech; several IPA heads decode the pronunciation and render the result with confidence-based colors. Confidence is based on model probability: - frame-level heads: top-1/top-2 margin aggregated over collapsed segments - ctc heads: mean probability over each emitted CTC run Suggested confidence thresholds (frame-level heads): ≥ 0.45 → green (confident) 0.25–0.45 → orange (medium) 0.12–0.25 → red (uncertain) < 0.12 → ? red bold (very uncertain, phone hidden) Display thresholds (ctc heads): > 0.76 → green 0.68–0.76 → yellow 0.41–0.68 → red ≤ 0.41 → ? """ import json import random import re import os import gradio as gr import librosa import nltk import torch from huggingface_hub import login as hf_login, snapshot_download from transformers import AutoConfig, AutoFeatureExtractor, AutoModel import numpy as np from ctc_decoding import ( decode_ctc_predictions, decode_ctc_predictions_beam, fold_phone_for_display, ) # HF_TOKEN is a private Space secret — never visible to demo users. # Required to avoid anonymous download throttling (downloads hang at 0 B). _HF_TOKEN = os.environ.get("HF_TOKEN") if _HF_TOKEN: hf_login(token=_HF_TOKEN, add_to_git_credential=False) # ── NLTK / CMU dict ─────────────────────────────────────────── nltk.download("cmudict", quiet=True) from nltk.corpus import cmudict # noqa: E402 CMU = cmudict.dict() # ── ARPAbet → IPA ───────────────────────────────────────────── _ARPA2IPA = { "B": "b", "D": "d", "G": "g", "P": "p", "T": "t", "K": "k", "CH": "tʃ", "JH": "dʒ", "F": "f", "V": "v", "TH": "θ", "DH": "ð", "S": "s", "Z": "z", "SH": "ʃ", "ZH": "ʒ", "HH": "h", "M": "m", "N": "n", "NG": "ŋ", "L": "l", "R": "r", "W": "w", "Y": "j", "IY": "i", "IH": "ɪ", "EH": "ɛ", "EY": "eɪ", "AE": "æ", "AA": "ɑ", "AW": "aʊ", "AY": "aɪ", "AH": "ə", "AO": "ɔ", "OY": "ɔɪ", "OW": "oʊ", "UH": "ʊ", "UW": "u", "ER": "ɝ", } _STRESSABLE_ARPA = { "IY", "IH", "EH", "EY", "AE", "AA", "AW", "AY", "AH", "AO", "OY", "OW", "UH", "UW", "ER", } THRESHOLD=0.5 def _arpa_phones_to_ipa(phones: list) -> str: """['HH', 'AH0', 'L', 'OW1'] → 'həloʊ'""" out = [] for p in phones: match = re.match(r"^([A-Z]+)(\d)?$", p.upper()) key = match.group(1) if match else re.sub(r"\d+$", "", p).upper() stress = match.group(2) if match else None ipa = _ARPA2IPA.get(key, p.lower()) if key in _STRESSABLE_ARPA: if stress == "1": ipa = "ˈ" + ipa elif stress == "2": ipa = "." + ipa out.append(ipa) return "".join(out) # ── Word list (top-5000 common words present in CMU dict) ───── print("Building word list…") try: from wordfreq import top_n_list _candidates = top_n_list("en", 30_000) except ImportError: # Fallback: use NLTK Brown corpus frequencies nltk.download("brown", quiet=True) from collections import Counter from nltk.corpus import brown _freq = Counter(w.lower() for w in brown.words() if w.isalpha()) _candidates = [w for w, _ in _freq.most_common(30_000)] WORDS = [w for w in _candidates if w.isalpha() and len(w) >= 3 and w in CMU][:5000] print(f" word list: {len(WORDS)} words") def _word_ipa(word: str) -> str: """Return IPA for first CMU pronunciation variant of *word*.""" phones = CMU[word][0] return "/" + _arpa_phones_to_ipa(phones) + "/" def get_random_word() -> tuple[str, str]: word = random.choice(WORDS) return word.upper(), _word_ipa(word) # ── Model ───────────────────────────────────────────────────── BACKBONE_ID = "utter-project/mHuBERT-147" HEAD_REPO = "istomin9192/mHuBERT-147-ipa-head" CTC_FT_LOCAL_DIR = "/cuda/pron/mHuBERT-147-ipa-ctc-ft" CTC_FT_REPO = "istomin9192/mHuBERT-147-ipa-ctc-ft" LINEAR_CTC_FT_LOCAL_DIR = "/cuda/pron/mHuBERT-147-ipa-linear-ctc-ft" LINEAR_CTC_FT_REPO = "istomin9192/mHuBERT-147-ipa-linear-ctc-ft" MAX_SAMPLES = 15 * 16_000 SR = 16_000 device = "cuda" if torch.cuda.is_available() else "cpu" print("Loading backbone…") head_repo_dir = snapshot_download(repo_id=HEAD_REPO, token=_HF_TOKEN) lstm_head_dir = os.path.join(head_repo_dir, "lstm_v1") phone_mask_head_dir = os.path.join(head_repo_dir, "phone_mask_v1") ctc_head_dir = os.path.join(head_repo_dir, "ctc_v1") conformer_head_dir = os.path.join(head_repo_dir, "conformer_v1") ipa_map_path = os.path.join(head_repo_dir, "ipa_map.json") ctc_ft_model_dir = CTC_FT_LOCAL_DIR if os.path.isdir(CTC_FT_LOCAL_DIR) else CTC_FT_REPO linear_ctc_ft_model_dir = ( LINEAR_CTC_FT_LOCAL_DIR if os.path.isdir(LINEAR_CTC_FT_LOCAL_DIR) else LINEAR_CTC_FT_REPO ) lstm_config = AutoConfig.from_pretrained(lstm_head_dir, trust_remote_code=True) phone_mask_config = AutoConfig.from_pretrained(phone_mask_head_dir, trust_remote_code=True) ctc_config = AutoConfig.from_pretrained(ctc_head_dir, trust_remote_code=True) conformer_config = AutoConfig.from_pretrained(conformer_head_dir, trust_remote_code=True) ctc_ft_config = AutoConfig.from_pretrained(ctc_ft_model_dir, trust_remote_code=True) linear_ctc_ft_config = AutoConfig.from_pretrained(linear_ctc_ft_model_dir, trust_remote_code=True) backbone_id = getattr(lstm_config, "base_model", BACKBONE_ID) feature_extractor = AutoFeatureExtractor.from_pretrained(backbone_id, token=_HF_TOKEN) backbone = AutoModel.from_pretrained(backbone_id, token=_HF_TOKEN).to(device).eval() ctc_ft_feature_extractor = AutoFeatureExtractor.from_pretrained( ctc_ft_model_dir, trust_remote_code=True, token=_HF_TOKEN, ) linear_ctc_ft_feature_extractor = AutoFeatureExtractor.from_pretrained( linear_ctc_ft_model_dir, trust_remote_code=True, token=_HF_TOKEN, ) print("Loading IPA heads…") lstm_head = AutoModel.from_pretrained( lstm_head_dir, trust_remote_code=True, ).to(device).eval() phone_mask_head = AutoModel.from_pretrained( phone_mask_head_dir, trust_remote_code=True, ).to(device).eval() ctc_head = AutoModel.from_pretrained( ctc_head_dir, trust_remote_code=True, ).to(device).eval() conformer_head = AutoModel.from_pretrained( conformer_head_dir, trust_remote_code=True, ).to(device).eval() ctc_ft_model = AutoModel.from_pretrained( ctc_ft_model_dir, trust_remote_code=True, token=_HF_TOKEN, ).to(device).eval() linear_ctc_ft_model = AutoModel.from_pretrained( linear_ctc_ft_model_dir, trust_remote_code=True, token=_HF_TOKEN, ).to(device).eval() with open(ipa_map_path, encoding="utf-8") as f: id2phone: dict = json.load(f)["id2phone"] # str(int) → phone def _prepare_audio(audio, sr): if audio is None or len(audio) == 0: return None wav = np.asarray(audio) if wav.ndim > 1: wav = wav.mean(axis=1) wav = wav.astype(np.float32) if sr != SR: wav = librosa.resample(wav, orig_sr=sr, target_sr=SR) wav = wav[:MAX_SAMPLES] return wav # ── Inference ───────────────────────────────────────────────── def _collapse_predictions(pred_ids: list[int], frame_conf: list[float]) -> tuple[list, list]: """Return (collapsed_display_phones, collapsed_display_conf_floats). Confidence is computed in the model's original raw class space as the margin between the top-1 and top-2 probabilities at each frame. Raw consecutive identical predictions are first collapsed into raw segments. For each raw segment, confidence is computed as the mean margin over the inner frames of the segment when possible (otherwise over all frames). Segment labels are then folded for display. If adjacent folded labels are identical, they are merged visually, and the merged confidence is the mean of the confidences of the merged raw segments. """ if not pred_ids or not frame_conf: return [], [] # 1) Collapse raw consecutive identical predictions into raw segments raw_seg_labels: list[str] = [] raw_seg_conf: list[float] = [] prev_pid = None run_conf: list[float] = [] for pid, c in zip(pred_ids, frame_conf): if pid != prev_pid: if prev_pid is not None: seg_conf = run_conf[1:-1] if len(run_conf) >= 3 else run_conf raw_seg_labels.append(id2phone[str(prev_pid)]) raw_seg_conf.append(sum(seg_conf) / len(seg_conf)) prev_pid = pid run_conf = [c] else: run_conf.append(c) if prev_pid is not None: seg_conf = run_conf[1:-1] if len(run_conf) >= 3 else run_conf raw_seg_labels.append(id2phone[str(prev_pid)]) raw_seg_conf.append(sum(seg_conf) / len(seg_conf)) # 2) Fold raw segment labels for display disp_seg_labels = [fold_phone_for_display(ph) for ph in raw_seg_labels] # 3) Merge adjacent identical folded labels visually collapsed: list[str] = [] conf_seg: list[float] = [] prev_ph = None run_seg_conf: list[float] = [] for ph, c in zip(disp_seg_labels, raw_seg_conf): if ph != prev_ph: if prev_ph is not None: collapsed.append(prev_ph) conf_seg.append(sum(run_seg_conf) / len(run_seg_conf)) prev_ph = ph run_seg_conf = [c] else: run_seg_conf.append(c) if prev_ph is not None: collapsed.append(prev_ph) conf_seg.append(sum(run_seg_conf) / len(run_seg_conf)) return collapsed, conf_seg @torch.no_grad() def _infer(audio, a_sr) -> dict[str, tuple[list, list]]: wav = _prepare_audio(audio, a_sr) if wav is None: return { "ctc_ft_v1": ([], []), "linear_ctc_ft_v1_beam": ([], []), "conformer_v1": ([], []), "ctc_v1": ([], []), "phone_mask_v1": ([], []), "lstm_v1": ([], []), } inp = feature_extractor(wav, sampling_rate=SR, return_tensors="pt") emb = backbone(inp.input_values.to(device)).last_hidden_state # [1,T,D] ctc_ft_inp = ctc_ft_feature_extractor(wav, sampling_rate=SR, return_tensors="pt") linear_ctc_ft_inp = linear_ctc_ft_feature_extractor(wav, sampling_rate=SR, return_tensors="pt") lstm_outputs = lstm_head(emb) lstm_probs = torch.softmax(lstm_outputs.logits[0], dim=-1) # [T,C] lstm_top2 = lstm_probs.topk(2, dim=-1).values lstm_conf = (lstm_top2[:, 0] - lstm_top2[:, 1]).tolist() lstm_pred_ids = lstm_probs.argmax(dim=-1).tolist() phone_mask_outputs = phone_mask_head(emb) phone_probs = torch.softmax(phone_mask_outputs.phone_logits[0], dim=-1) phone_top2 = phone_probs.topk(2, dim=-1).values phone_conf = phone_top2[:, 0] - phone_top2[:, 1] phone_pred_ids = phone_probs.argmax(dim=-1) utility = torch.sigmoid(phone_mask_outputs.utility_logits[0]) mask = utility > THRESHOLD masked_pred_ids = phone_pred_ids[mask].tolist() masked_conf = phone_conf[mask].tolist() ctc_outputs = ctc_head(emb) ctc_probs = torch.softmax(ctc_outputs.logits[0], dim=-1) ctc_blank_id = ctc_config.architecture["blank_id"] conformer_outputs = conformer_head(emb, input_lengths=[emb.shape[1]]) conformer_probs = torch.softmax(conformer_outputs.logits[0], dim=-1) conformer_blank_id = conformer_config.architecture["blank_id"] ctc_ft_outputs = ctc_ft_model(input_values=ctc_ft_inp.input_values.to(device)) ctc_ft_probs = torch.softmax(ctc_ft_outputs.logits[0], dim=-1) ctc_ft_blank_id = ctc_ft_config.architecture["blank_id"] linear_ctc_ft_outputs = linear_ctc_ft_model(input_values=linear_ctc_ft_inp.input_values.to(device)) linear_ctc_ft_probs = torch.softmax(linear_ctc_ft_outputs.logits[0], dim=-1) linear_ctc_ft_blank_id = linear_ctc_ft_config.architecture["blank_id"] return { "ctc_ft_v1": decode_ctc_predictions(ctc_ft_probs, ctc_ft_blank_id, id2phone), "linear_ctc_ft_v1_beam": decode_ctc_predictions_beam( linear_ctc_ft_probs, linear_ctc_ft_blank_id, id2phone, beam_size=8, apply_cambridge_norm=True, ), "conformer_v1": decode_ctc_predictions(conformer_probs, conformer_blank_id, id2phone), "lstm_v1": _collapse_predictions(lstm_pred_ids, lstm_conf), "phone_mask_v1": _collapse_predictions(masked_pred_ids, masked_conf), "ctc_v1": decode_ctc_predictions(ctc_probs, ctc_blank_id, id2phone), } # ── Colored HTML rendering ──────────────────────────────────── _FONT_STYLE = ( "font-family:'Charis SIL',serif,sans-serif;" "font-size:1.6rem;letter-spacing:0.06em" ) def _phone_span_frame(phone: str, conf: float) -> str: """Color scheme for frame-level heads based on margin confidence.""" if phone == "sil": return " " esc = phone.replace("&", "&").replace("<", "<").replace(">", ">") if conf < 0.12: return '?' elif conf < 0.25: return f'{esc}' elif conf < 0.45: return f'{esc}' else: return f'{esc}' def _phone_span_ctc(phone: str, conf: float) -> str: """Color scheme for ctc_v1 based on phone posterior.""" if phone == "sil": return " " esc = phone.replace("&", "&").replace("<", "<").replace(">", ">") if conf <= 0.41: return '?' elif conf <= 0.68: return f'{esc}' elif conf <= 0.76: return f'{esc}' else: return f'{esc}' def render_colored_ipa(collapsed: list, conf_list: list, mode: str = "frame") -> str: span_fn = _phone_span_ctc if mode == "ctc" else _phone_span_frame inner = "".join(span_fn(ph, c) for ph, c in zip(collapsed, conf_list)) return ( f'
' f"[{inner.strip()}]" f"
" ) def render_model_result(title: str, collapsed: list, conf_list: list, mode: str = "frame") -> str: if not collapsed: body = '
No phones decoded.
' else: body = render_colored_ipa(collapsed, conf_list, mode=mode) return ( '
' f'
{title}
' f"{body}" "
" ) # ── Gradio callbacks ────────────────────────────────────────── def cb_next_word(): word, ipa = get_random_word() return word, ipa, "", None # word, ref_ipa, predicted_html, audio (cleared) def cb_lookup_ipa(word: str) -> str: """Look up IPA for a user-typed word; return empty string if not found.""" w = word.strip().lower() if not w: return "" if w not in CMU: return "not in dict" return _word_ipa(w) def cb_recognize(audio, show_all_models: bool) -> str: if audio is None: return "" a_sr, wav = audio results = _infer(wav, a_sr) if not any(collapsed for collapsed, _ in results.values()): return "" for model_name, (collapsed, conf) in results.items(): if collapsed: scores = " ".join(str(int(c * 100)) for c in conf) print(f"Recognized [{model_name}]:", " ".join(collapsed), "/", scores) else: print(f"Recognized [{model_name}]: no phones decoded") html = ( render_model_result("conformer_ctc_v1", *results["conformer_v1"], mode="ctc") + render_model_result("linear_ctc_ft_v1_beam", *results["linear_ctc_ft_v1_beam"], mode="ctc") ) if show_all_models: html += ( render_model_result("ctc_ft_v1", *results["ctc_ft_v1"], mode="ctc") + render_model_result("ctc_v1", *results["ctc_v1"], mode="ctc") + render_model_result("phone_mask_v1", *results["phone_mask_v1"], mode="frame") + render_model_result("lstm_v1", *results["lstm_v1"], mode="frame") ) return html def cb_toggle_recognize(audio): return gr.update(interactive=audio is not None) def cb_trim_audio(audio): if audio is None: return None a_sr, wav = audio wav = _prepare_audio(wav, a_sr) if wav is None: return None wav_trimmed, idx = librosa.effects.trim(wav, top_db=30) if len(wav_trimmed) == 0: return None start, end = idx pad = int(0.16 * SR) start = max(0, start - pad) end = min(len(wav), end + pad) wav_trimmed = wav[start:end].copy() if len(wav_trimmed) == 0: return None # Smooth the cut boundaries to avoid audible clicks after trimming. fade_samples = min(int(0.04 * SR), len(wav_trimmed) // 2) if fade_samples > 1: fade_in = np.linspace(0.0, 1.0, fade_samples, dtype=np.float32) fade_out = fade_in[::-1] wav_trimmed[:fade_samples] *= fade_in wav_trimmed[-fade_samples:] *= fade_out wav_trimmed = wav_trimmed.astype(np.int16) return (SR, wav_trimmed) def cb_trim_and_recognize(audio, show_all_models: bool): trimmed_audio = cb_trim_audio(audio) recognize_state = cb_toggle_recognize(trimmed_audio) predicted_html = cb_recognize(trimmed_audio, show_all_models) if trimmed_audio is not None else "" return trimmed_audio, predicted_html, recognize_state # ── UI ──────────────────────────────────────────────────────── css = """ @import url('https://fonts.googleapis.com/css2?family=Charis+SIL&family=Open+Sans&display=swap'); #page { max-width: 860px; margin: auto; } #ref_ipa textarea { font-family:'Charis SIL',serif; font-size:1.3rem; letter-spacing:0.06em; text-align:center; } #word_display textarea { font-family:'Open Sans',sans-serif; font-size:1.1rem; letter-spacing:0.02em; font-weight:400; text-align:center; } """ theme = gr.themes.Soft(primary_hue=gr.themes.utils.colors.indigo).set( button_primary_background_fill="#4F46E5", button_primary_background_fill_hover="#3730A3", button_primary_text_color="white", ) with gr.Blocks(theme=theme, css=css) as demo: with gr.Column(elem_id="page"): gr.Markdown( """
# 🗣️ Pronunciation Checker A random word is shown with its IPA transcription. Record yourself saying it — the model will show your pronunciation colored by confidence.
""" ) with gr.Row(): word_display = gr.Textbox( label="Word", interactive=True, text_align="center", placeholder="Type a word…", elem_id="word_display", ) ref_ipa = gr.Textbox( label="Reference IPA", interactive=False, elem_id="ref_ipa", text_align="center", ) next_btn = gr.Button("🔀 New word", variant="secondary") # with gr.Row(): audio_input = gr.Audio(type="numpy", label="🎤 Record or upload") predicted = gr.HTML(label="Predicted IPA") recognize_btn = gr.Button("🔍 Recognize", variant="primary", interactive=False) show_all_models = gr.Checkbox( label="Show all model variants", value=False, ) # ── events next_btn.click( cb_next_word, outputs=[word_display, ref_ipa, predicted, audio_input], ) recognize_btn.click(cb_recognize, inputs=[audio_input, show_all_models], outputs=predicted) audio_input.stop_recording( cb_trim_and_recognize, inputs=[audio_input, show_all_models], outputs=[audio_input, predicted, recognize_btn], ) audio_input.change(cb_toggle_recognize, inputs=audio_input, outputs=recognize_btn) # ── user types their own word → look up IPA on Enter or blur for _ev in (word_display.submit, word_display.blur): _ev(cb_lookup_ipa, inputs=word_display, outputs=ref_ipa) # ── on load: show first word demo.load(cb_next_word, outputs=[word_display, ref_ipa, predicted, audio_input]) demo.launch()