Deploy FALCON demo (app + bundled MFA G2P assets + example inputs)

.gitattributes

@@ -1,35 +1,2 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.fst filter=lfs diff=lfs merge=lfs -text
+*.wav filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,31 @@
 ---
-title: FALCON
-emoji: 📊
-colorFrom: green
-colorTo: gray
+title: FALCON Forced Aligner
+emoji: 🦅
+colorFrom: indigo
+colorTo: blue
 sdk: gradio
-sdk_version: 6.19.0
-python_version: '3.13'
+sdk_version: 4.44.1
+python_version: "3.8"
 app_file: app.py
 pinned: false
+license: mit
+short_description: Neural forced alignment via Soft Dynamic Programming
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# FALCON — Forced Alignment through Contrastive Optimization Networks
+
+Interactive demo of **FALCON**, a fully differentiable neural forced aligner that predicts
+precise **phoneme- and word-level** boundary timestamps from a waveform + transcript, using a
+Soft Dynamic Programming decoder.
+
+Upload audio + a transcript (`.phn` / `.wrd` / `.txt`), choose the options, and get a boundary
+table, a downloadable Praat `.TextGrid`, and a time-aligned visualization (waveform ·
+spectrogram · phoneme posteriors · Soft-DP path · contrastive score).
+
+- **Paper:** *Fully Differentiable Neural Forced Alignment via Soft Dynamic Programming* — [arXiv:2606.25460](https://arxiv.org/abs/2606.25460)
+- **Code:** https://github.com/MLSpeech/FALCON
+- **Weights:** https://huggingface.co/MLSpeech/FALCON-weights
+
+Example inputs are in `assets/` — the TIMIT sentence *"Don't ask me to carry an oily rag like that."*
+in every supported format. The checkpoints are downloaded automatically from the weights repo on
+first use (this runs on a free CPU Space, so the first alignment takes a moment to fetch a model).

app.py

@@ -0,0 +1,528 @@
+"""
+FALCON web demo — interactive forced alignment in the browser.
+
+Two pretrained checkpoints are used by FALCON (under pretrained_models/):
+  - falcon_timit_english.pt       — TIMIT-trained, best for English phoneme alignment
+  - falcon_joint_multilingual.pt  — joint TIMIT+Buckeye model; best for cross-lingual /
+                                   multilingual zero-shot alignment (Dutch, German,
+                                   Hebrew, ...) at both phoneme and word level.
+
+The app picks one automatically from the `Language` radio (english → TIMIT,
+multilingual → joint); a custom .pt upload overrides both.
+
+For HuggingFace Spaces deployment, set Space Secrets:
+  HF_MODEL_REPO   — e.g. "MLSpeech/FALCON-weights"
+  HF_TOKEN        — only needed for private repos
+The app will download `falcon_timit_english.pt` and `falcon_joint_multilingual.pt`
+from that repo on first use.
+"""
+import os
+import re
+import shutil
+import sys
+import tempfile
+import threading
+import time
+
+import gradio as gr
+import textgrid
+import torchaudio
+
+import utils
+from predict import main_predict
+
+# On HF Spaces, point the "MFA-like" word G2P at the bundled dictionaries / G2P FST
+# and at this interpreter (which has pynini), so it works without a separate MFA
+# aligner conda env. No effect off Spaces — your local MFA install is used as-is.
+if os.environ.get("SPACE_ID"):
+    _SPACE_DIR = os.path.dirname(os.path.abspath(__file__))
+    os.environ.setdefault("MFA_ROOT_DIR", os.path.join(_SPACE_DIR, "mfa_assets"))
+    os.environ.setdefault("FDNFA_MFA_ENV_PY", sys.executable)
+
+# ── Checkpoint configuration ──────────────────────────────────────────────────
+
+SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+PRETRAINED_DIR = os.path.join(SCRIPT_DIR, "pretrained_models")
+
+CKPT_FILES = {
+    "english":      "falcon_timit_english.pt",
+    "buckeye":      "falcon_buckeye_english.pt",
+    "multilingual": "falcon_joint_multilingual.pt",
+}
+CKPT_LABELS = {
+    "english":      "Read English (recommended) — TIMIT model",
+    "buckeye":      "Spontaneous English (recommended) — Buckeye model",
+    "multilingual": "Multilingual — joint TIMIT+Buckeye model",
+}
+
+_ckpt_cache = {}
+
+def _resolve_ckpt(key: str):
+    """Resolve checkpoint path: cache → local file → HF Hub. Returns None if all fail."""
+    if key in _ckpt_cache and os.path.exists(_ckpt_cache[key]):
+        return _ckpt_cache[key]
+
+    filename = CKPT_FILES[key]
+    local_path = os.path.join(PRETRAINED_DIR, filename)
+    if os.path.exists(local_path):
+        _ckpt_cache[key] = local_path
+        return local_path
+
+    repo = os.environ.get("HF_MODEL_REPO", "")
+    if repo:
+        try:
+            from huggingface_hub import hf_hub_download
+            path = hf_hub_download(
+                repo_id=repo,
+                filename=filename,
+                token=os.environ.get("HF_TOKEN"),
+            )
+            _ckpt_cache[key] = path
+            return path
+        except Exception as exc:
+            print(f"[FALCON] HF Hub fetch failed for {filename}: {exc}")
+
+    return None
+
+_inference_lock = threading.Lock()
+
+# ── Heartbeat-based auto-shutdown (local runs only) ───────────────────────────
+# The open browser tab pings _heartbeat() every few seconds. A watchdog thread
+# exits the process when the pings stop (tab closed / browser crashed / unload
+# event dropped). _last_ping stays None until the first browser connects, so the
+# server never self-exits before anyone opens it.
+_last_ping = [None]
+_HEARTBEAT_TIMEOUT = 90  # secs of silence before the local server self-exits
+
+def _heartbeat():
+    _last_ping[0] = time.time()
+
+def _start_shutdown_watchdog():
+    def _watch():
+        while True:
+            time.sleep(5)
+            last = _last_ping[0]
+            if last is not None and (time.time() - last) > _HEARTBEAT_TIMEOUT:
+                os._exit(0)
+    threading.Thread(target=_watch, daemon=True).start()
+
+# ── Internal language routing ────────────────────────────────────────────────
+
+def _internal_language(lang: str, mode: str, ann_ext: str) -> str:
+    """
+    Map UI choices to the internal `language` flag understood by main_predict.
+
+      'english' = no G2P; assumes labels are already TIMIT-39 phonemes.
+      'dutch'   = G2P pipeline (panphon-based articulatory mapping). Used for:
+                    • any non-English language
+                    • word-level alignment (.wrd, words need phoneme decomposition)
+                    • plain text input (could be words or arbitrary phonemes)
+
+    NOTE: `ann_ext` here must be the *original* extension supplied by the user —
+    not the post-rewrite extension after a .txt → dummy .phn synthesis.
+    """
+    if lang == "english" and mode == "phoneme" and ann_ext.lower() == "phn":
+        return "english"
+    return "dutch"
+
+# ── Word-level G2P selection ──────────────────────────────────────────────────
+
+# Optional input-language hint -> (espeak voice, MFA voice or None). MFA ships
+# pronunciation models only for en/de/nl; everything else uses espeak, or "none"
+# (romanized characters -> LH39) when even espeak has no voice.
+G2P_LANG_CHOICES = [
+    "English (default)", "German", "Dutch", "Hebrew", "French",
+    "Spanish", "Italian", "Russian", "Portuguese", "Other / unknown",
+]
+_G2P_LANG_MAP = {
+    "English (default)": ("en-us", "en-us"),
+    "German":            ("de", "de"),
+    "Dutch":             ("nl", "nl"),
+    "Hebrew":            ("he", None),
+    "French":            ("fr", None),
+    "Spanish":           ("es", None),
+    "Italian":           ("it", None),
+    "Russian":           ("ru", None),
+    "Portuguese":        ("pt", None),
+    "Other / unknown":   ("en-us", None),
+}
+
+def _resolve_g2p(g2p_choice, lang_choice):
+    """Map the (G2P option, input-language) UI choices to a concrete backend.
+
+    Returns (backend, voice, note). backend in {"mfa", "espeak", "char"}. Honors
+    an explicit espeak / MFA-like / none choice but auto-falls-back when the
+    chosen backend has no model for the language; "Auto" picks the best available.
+    """
+    espeak_voice, mfa_voice = _G2P_LANG_MAP.get(lang_choice, ("en-us", "en-us"))
+    try:
+        import mfa_g2p
+        mfa_ok = mfa_voice is not None and mfa_g2p.mfa_available(mfa_voice)
+    except Exception:
+        mfa_ok = False
+    choice = (g2p_choice or "Auto").lower()
+
+    if choice.startswith("none"):
+        return "char", espeak_voice or "en-us", "none (romanization)"
+    if choice.startswith("mfa"):
+        if mfa_ok:
+            return "mfa", mfa_voice, "MFA-like"
+        if espeak_voice:
+            return "espeak", espeak_voice, "espeak (no MFA model for this language)"
+        return "char", "en-us", "none (no MFA/espeak model)"
+    if choice.startswith("espeak"):
+        if espeak_voice:
+            return "espeak", espeak_voice, "espeak"
+        if mfa_ok:
+            return "mfa", mfa_voice, "MFA-like (no espeak voice for this language)"
+        return "char", "en-us", "none"
+    # Auto (recommended)
+    if mfa_ok:
+        return "mfa", mfa_voice, "MFA-like (auto)"
+    if espeak_voice:
+        return "espeak", espeak_voice, "espeak (auto)"
+    return "char", "en-us", "none (auto)"
+
+# ── Core handler ──────────────────────────────────────────────────────────────
+
+OUTPUTS_NONE = (None, None, None, None, None)  # 5 None for the non-status outputs
+
+def run_alignment(audio_file, annotation_file, ckpt_upload, mode, lang,
+                  pretrained_choice, w_phi, g2p_choice="Auto (recommended)",
+                  lang_choice="English (default)",
+                  progress=gr.Progress(track_tqdm=True)):
+    if not audio_file or not annotation_file:
+        return ("Please upload both an audio file and an annotation file.", *OUTPUTS_NONE)
+
+    ckpt_to_use = ckpt_upload if ckpt_upload else _resolve_ckpt(pretrained_choice)
+    if not ckpt_to_use or not os.path.exists(ckpt_to_use):
+        return (
+            f"No checkpoint found. Expected {CKPT_FILES[pretrained_choice]} "
+            f"in {PRETRAINED_DIR}, or HF_MODEL_REPO set, or upload a .pt file.",
+            *OUTPUTS_NONE,
+        )
+
+    progress(0.1, desc="Preparing workspace...")
+    workspace = tempfile.mkdtemp(prefix="falcon_")
+    base = "input"
+    wav_path = os.path.join(workspace, f"{base}.wav")
+
+    original_ext = os.path.basename(annotation_file).split(".")[-1].lower()
+    ann_ext = original_ext
+    ann_path = os.path.join(workspace, f"{base}.{ann_ext}")
+
+    # Resample audio to 16 kHz mono
+    try:
+        audio, sr = torchaudio.load(audio_file)
+        if audio.shape[0] > 1:
+            audio = audio.mean(dim=0, keepdim=True)
+        if sr != 16000:
+            audio = torchaudio.functional.resample(audio, sr, 16000)
+        torchaudio.save(wav_path, audio, 16000)
+    except Exception as exc:
+        return (f"Audio error: {exc}", *OUTPUTS_NONE)
+
+    shutil.copy(annotation_file, ann_path)
+
+    # Capture the original input tokens (whatever the user supplied per line):
+    #   .phn → phoneme labels
+    #   .wrd → word labels
+    #   .txt → space-separated tokens (words or phonemes)
+    if original_ext == "txt":
+        with open(ann_path) as f:
+            orig_tokens = re.sub(r"[^\w\s]", "", f.read().strip()).split()
+        # TIMIT .txt files are "<start_sample> <end_sample> <sentence>" — drop the
+        # leading sample indices so they aren't mistaken for words.
+        if len(orig_tokens) >= 3 and orig_tokens[0].isdigit() and orig_tokens[1].isdigit():
+            orig_tokens = orig_tokens[2:]
+        if not orig_tokens:
+            return ("Text annotation is empty after stripping punctuation.", *OUTPUTS_NONE)
+        # Synthesize a uniform-segments dummy .phn so downstream code has timestamps.
+        audio_len = audio.shape[1]
+        interval = audio_len / len(orig_tokens)
+        ann_ext = "phn"
+        ann_path = os.path.join(workspace, f"{base}.{ann_ext}")
+        with open(ann_path, "w") as f:
+            for i, tok in enumerate(orig_tokens):
+                f.write(f"{int(i * interval)} {int((i + 1) * interval)} {tok}\n")
+    else:
+        with open(ann_path) as f:
+            orig_tokens = [ln.strip().split()[-1] for ln in f if ln.strip()]
+
+    # Route by ORIGINAL extension (post-rewrite ann_ext is "phn" for txt inputs).
+    language = _internal_language(lang, mode, original_ext)
+
+    # Word-level G2P: convert orthographic words -> LH39 phonemes with the chosen
+    # front-end (espeak, or the MFA english_us_arpa G2P used in the paper), then
+    # align via the stock phoneme path. Replaces the legacy letter-by-letter
+    # mapping. Only applies to real word input (.wrd / .txt / .word).
+    # Word/text inputs (.wrd / .txt) always go through the proper word -> LH39 G2P,
+    # regardless of the phoneme/word toggle — otherwise phoneme mode would fall back
+    # to a crude character mapping and misalign. (.phn is the phoneme-input path.)
+    app_mapped_ph = None
+    word_g2p_note = ""
+    if original_ext in ("wrd", "txt", "word") and orig_tokens:
+        import word_g2p
+        backend, voice, g2p_note = _resolve_g2p(g2p_choice, lang_choice)
+        try:
+            app_mapped_ph = [word_g2p.word_to_lh39(tok, voice=voice, backend=backend)
+                             for tok in orig_tokens]
+        except Exception as g2p_exc:
+            # A missing model must never break the whole run — fall back.
+            print(f"[FALCON] G2P backend '{backend}' failed ({g2p_exc}); falling back.")
+            try:
+                app_mapped_ph = [word_g2p.word_to_lh39(tok, voice=voice or "en-us",
+                                                       backend="espeak")
+                                 for tok in orig_tokens]
+                g2p_note += " → espeak fallback"
+            except Exception:
+                app_mapped_ph = [word_g2p.word_to_lh39(tok, backend="char")
+                                 for tok in orig_tokens]
+                g2p_note += " → none fallback"
+        word_g2p_note = f" Word G2P: {g2p_note}."
+        phons = [ph for seq in app_mapped_ph for ph in seq] or ["sil"]
+        # Rewrite the annotation as a dummy uniform-time .phn of LH39 phonemes so
+        # the stock English/phoneme aligner runs on them (no further G2P).
+        ann_ext = "phn"
+        ann_path = os.path.join(workspace, f"{base}.{ann_ext}")
+        audio_len = audio.shape[1]
+        interval = audio_len / max(1, len(phons))
+        with open(ann_path, "w") as f:
+            for i, ph in enumerate(phons):
+                f.write(f"{int(i * interval)} {int((i + 1) * interval)} {ph}\n")
+        language = "english"   # phonemes are already LH39; skip the internal G2P
+
+    progress(0.3, desc="Running alignment...")
+    try:
+        with _inference_lock:
+            utils.set_dp_matrix_out_dir(workspace)
+            pred_bound, _truth_bound, mapped_ph = main_predict(
+                wav=wav_path,
+                ckpt=ckpt_to_use,
+                w_phi=w_phi,
+                language=language,
+                annotation=ann_ext,
+            )
+            utils.set_dp_matrix_out_dir(None)
+            progress(0.6, desc="Rendering aligned visualization...")
+            # Time-aligned representations as one stacked figure (waveform,
+            # spectrogram, phoneme posteriors, Soft-DP matrix + path, contrastive
+            # score) — all on the same time axis with predicted boundaries overlaid.
+            panels_path = os.path.join(workspace, "panels.png")
+            try:
+                import falcon_viz
+                falcon_viz.make_alignment_panels(
+                    wav=wav_path, ckpt=ckpt_to_use, out_path=panels_path,
+                    language=language, annotation=ann_ext,
+                    show_truth=(original_ext == "phn" and mode == "phoneme"),
+                )
+            except Exception as viz_exc:
+                print(f"[FALCON] panel viz failed: {viz_exc}")
+                panels_path = None
+    except Exception as exc:
+        utils.set_dp_matrix_out_dir(None)
+        return (f"Inference error: {exc}", *OUTPUTS_NONE)
+
+    # When the app did the word-level G2P itself, use its per-word LH39 phoneme
+    # lists for the two-table / TextGrid word tier (main_predict's English path
+    # returns mapped_ph=None).
+    if app_mapped_ph is not None:
+        mapped_ph = app_mapped_ph
+
+    progress(0.8, desc="Building outputs...")
+
+    pred_bound_list = [float(t) for t in pred_bound]
+
+    # ── Build two tables ────────────────────────���────────────────────────────
+    # 1) LH39 phonemes  — the aligner's direct output, one row per pred_bound.
+    # 2) Original tokens — words (.wrd / .txt) or non-LH39 phonemes (multilingual
+    #    .phn). Only populated when the G2P path was used (mapped_ph != None);
+    #    for english+phoneme+.phn the LH39 phonemes ARE the original, so the
+    #    second table is left empty.
+    # Every input token gets a row in the table even if its predicted interval
+    # is degenerate; degenerate intervals are still kept out of the TextGrid
+    # (Praat rejects zero-length).
+
+    if mapped_ph is not None:
+        phn_labels = [ph for seq in mapped_ph for ph in seq]
+    else:
+        phn_labels = orig_tokens
+
+    table_phonemes, phn_intervals = [], []
+    t0 = 0.0
+    for i, t1 in enumerate(pred_bound_list):
+        lbl = phn_labels[i] if i < len(phn_labels) else ""
+        table_phonemes.append([round(t0, 3), round(t1, 3), lbl])
+        if t1 > t0:
+            phn_intervals.append((t0, t1, lbl))
+        t0 = t1
+
+    table_original, orig_intervals = [], []
+    if mapped_ph is not None:
+        counts_per_token = [len(seq) for seq in mapped_ph]
+        cumulative = 0
+        t0 = 0.0
+        for i, count in enumerate(counts_per_token):
+            cumulative += count
+            if cumulative - 1 >= len(pred_bound_list):
+                break
+            t1 = pred_bound_list[cumulative - 1]
+            lbl = orig_tokens[i] if i < len(orig_tokens) else ""
+            table_original.append([round(t0, 3), round(t1, 3), lbl])
+            if t1 > t0:
+                orig_intervals.append((t0, t1, lbl))
+            t0 = t1
+
+    # ── TextGrid: phones tier always, original tier when applicable ──────────
+    max_time = phn_intervals[-1][1] if phn_intervals else 0.0
+    tg = textgrid.TextGrid(minTime=0, maxTime=max_time)
+    tier_phn = textgrid.IntervalTier(name="phones", minTime=0, maxTime=max_time)
+    for t0_iv, t1_iv, lbl_iv in phn_intervals:
+        tier_phn.add(minTime=t0_iv, maxTime=t1_iv, mark=lbl_iv)
+    tg.append(tier_phn)
+    if orig_intervals:
+        # Tier name reflects what the original layer represents.
+        if mode == "word":
+            orig_tier_name = "words"
+        elif original_ext == "phn":
+            orig_tier_name = "phones_original"
+        else:
+            orig_tier_name = "tokens"
+        tier_orig = textgrid.IntervalTier(name=orig_tier_name, minTime=0, maxTime=max_time)
+        for t0_iv, t1_iv, lbl_iv in orig_intervals:
+            tier_orig.add(minTime=t0_iv, maxTime=t1_iv, mark=lbl_iv)
+        tg.append(tier_orig)
+    tg_path = os.path.join(workspace, f"{base}.TextGrid")
+    tg.write(tg_path)
+
+    # Status note: the .phn-as-word case produces a second "words" table that
+    # actually contains phonemes — surface this in the status so it's not
+    # mistaken for a bug.
+    status_note = word_g2p_note
+    if mode == "word" and original_ext == "phn":
+        status_note += (" Note: input was phoneme-level (.phn) but mode=word "
+                        "— the 'original' table shows input phonemes since no "
+                        "word annotations were provided.")
+
+    return (
+        "Done." + status_note,
+        audio_file,
+        panels_path if panels_path and os.path.exists(panels_path) else None,
+        tg_path,
+        table_phonemes,
+        table_original,
+    )
+
+# ── UI ────────────────────────────────────────────────────────────────────────
+
+PRETRAINED_RADIO_CHOICES = [
+    (CKPT_LABELS["english"],      "english"),
+    (CKPT_LABELS["buckeye"],      "buckeye"),
+    (CKPT_LABELS["multilingual"], "multilingual"),
+]
+
+with gr.Blocks(title="FALCON Forced Aligner", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# FALCON: Forced Alignment through Contrastive Optimization Networks")
+    gr.Markdown(
+        "Upload a speech file and a transcript to predict precise phoneme or word boundaries "
+        "using Soft Dynamic Programming."
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.Markdown("### Inputs")
+            audio_in  = gr.Audio(label="Audio file (any sample rate)", type="filepath")
+            ann_in    = gr.File(label="Annotation (.phn / .wrd / .txt)")
+            mode_in   = gr.Radio(["phoneme", "word"], value="phoneme", label="Mode")
+            lang_in   = gr.Radio(["english", "multilingual"], value="english", label="Language")
+            # Word-level G2P front-end (word mode only). The espeak voice / MFA
+            # dictionary is chosen automatically from the optional input-language
+            # hint below; "Auto" also picks the best available backend.
+            g2p_in    = gr.Radio(
+                ["Auto (recommended)", "espeak", "MFA-like",
+                 "none — romanization (not recommended; only for languages with no G2P model)"],
+                value="Auto (recommended)",
+                label="Word G2P (used only in word mode)",
+            )
+            lang_g2p_in = gr.Dropdown(
+                G2P_LANG_CHOICES,
+                value="English (default)",
+                label="Input language (optional, recommended — improves G2P choice)",
+            )
+            pretrained_in = gr.Radio(
+                choices=PRETRAINED_RADIO_CHOICES,
+                value="english",
+                label="Pretrained checkpoint (auto-follows Language; override here if you want)",
+            )
+            ckpt_in   = gr.File(label="Or upload a custom checkpoint (.pt) — overrides pretrained",
+                                file_types=[".pt"], type="filepath")
+            wphi_in   = gr.Slider(0.0, 1.0, value=0.5, step=0.01,
+                                  label="φ weight (acoustic ↔ linguistic)")
+            btn       = gr.Button("Run Alignment", variant="primary")
+            status    = gr.Textbox(label="Status", interactive=False)
+
+        with gr.Column(scale=2):
+            gr.Markdown("### Outputs")
+            with gr.Tabs():
+                with gr.Tab("Alignment Data"):
+                    audio_out  = gr.Audio(label="Playback", interactive=False)
+                    table_phn_out = gr.Dataframe(
+                        headers=["Start (s)", "End (s)", "Phoneme (LH39)"],
+                        label="LH39 phonemes — aligner output",
+                    )
+                    table_orig_out = gr.Dataframe(
+                        headers=["Start (s)", "End (s)", "Label"],
+                        label="Original input layer (words / non-LH39 phonemes) — empty for English phoneme alignment",
+                    )
+                    tg_out     = gr.File(label="Download TextGrid (carries both tiers when applicable)")
+                with gr.Tab("Visualizations"):
+                    img_panels = gr.Image(
+                        label="Time-aligned representations — waveform · spectrogram · phoneme posteriors · Soft-DP path · contrastive score (shared time axis; predicted boundaries overlaid). Click to enlarge.",
+                        show_download_button=True,
+                    )
+
+    # Auto-flip the pretrained-checkpoint radio when the user changes language.
+    lang_in.change(fn=lambda v: v, inputs=lang_in, outputs=pretrained_in)
+
+    btn.click(
+        fn=run_alignment,
+        inputs=[audio_in, ann_in, ckpt_in, mode_in, lang_in, pretrained_in, wphi_in,
+                g2p_in, lang_g2p_in],
+        outputs=[status, audio_out, img_panels,
+                 tg_out, table_phn_out, table_orig_out],
+    )
+
+    # Auto-shutdown the local Python server when the user closes the browser
+    # tab. Disabled on HuggingFace Spaces (where SPACE_ID is set automatically)
+    # because the container is shared across visitors — one tab close should
+    # not tear down everyone else's session.
+    if not os.environ.get("SPACE_ID"):
+        # Fast path: unload events click a hidden Shutdown button -> os._exit.
+        shutdown_btn = gr.Button("Shutdown", visible=False, elem_id="falcon-shutdown-btn")
+        shutdown_btn.click(fn=lambda: os._exit(0), inputs=[], outputs=[])
+
+        # Guaranteed fallback: the page heartbeats; the watchdog exits if it stops.
+        hb_btn = gr.Button("hb", visible=False, elem_id="falcon-heartbeat-btn")
+        hb_btn.click(fn=_heartbeat, inputs=[], outputs=[],
+                     show_progress="hidden", queue=False)
+        _start_shutdown_watchdog()
+
+        demo.load(None, None, None, js="""
+            () => {
+                const stop = () => {
+                    const btn = document.getElementById('falcon-shutdown-btn');
+                    if (btn) btn.click();
+                };
+                window.addEventListener('beforeunload', stop);
+                window.addEventListener('pagehide', stop);
+                const beat = () => {
+                    const hb = document.getElementById('falcon-heartbeat-btn');
+                    if (hb) hb.click();
+                };
+                beat();
+                setInterval(beat, 10000);
+            }
+        """)
+
+if __name__ == "__main__":
+    demo.launch()

assets/fasw0sa2.phn

@@ -0,0 +1,36 @@
+0 2200 h#
+2200 2510 d
+2510 5229 ow
+5229 5576 n
+5576 7000 q
+7000 9460 ae
+9460 10480 s
+10480 11760 epi
+11760 12640 m
+12640 14720 iy
+14720 15550 tcl
+15550 16000 t
+16000 16760 ix
+16760 18070 kcl
+18070 19050 k
+19050 20804 eh
+20804 22099 r
+22099 23979 iy
+23979 25560 eh
+25560 26223 n
+26223 26674 ax
+26674 27897 q
+27897 30161 ao
+30161 31080 l
+31080 32292 iy
+32292 34339 r
+34339 38440 ae
+38440 39160 gcl
+39160 39530 g
+39530 40393 l
+40393 42120 ay
+42120 44040 kcl
+44040 44864 dh
+44864 48680 ae
+48680 49660 tcl
+49660 50640 h#

assets/fasw0sa2.txt

@@ -0,0 +1 @@
+0 50688 Don't ask me to carry an oily rag like that.

assets/fasw0sa2.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:32a6bfe7791589a1467ee7993de348bf249ef05186fe97e215d379db85670f7e
+size 101420

assets/fasw0sa2.wrd

@@ -0,0 +1,10 @@
+2200 7000 don't
+7000 10480 ask
+11760 14720 me
+14720 16760 to
+16760 23979 carry
+23979 26223 an
+26223 32292 oily
+32292 39530 rag
+39530 44040 like
+44040 49660 that

dataloader.py

@@ -0,0 +1,141 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, Dataset
+torch.multiprocessing.set_sharing_strategy('file_system')
+from tqdm import tqdm
+import numpy as np
+import os
+from os.path import join, basename
+from boltons.fileutils import iter_find_files
+import soundfile as sf
+import librosa
+import pickle
+from multiprocessing import Pool
+import random
+import torchaudio
+import math
+from torchaudio.datasets import LIBRISPEECH
+
+
+def collate_fn_padd(batch):
+    """collate_fn_padd
+    Padds batch of variable length
+
+    :param batch:
+    """
+    # get sequence lengths
+    spects = [t[0] for t in batch]
+    segs = [t[1] for t in batch]
+    labels = [t[2] for t in batch]
+    lengths = [t[3] for t in batch]
+    fnames = [t[4] for t in batch]
+
+    padded_spects = torch.nn.utils.rnn.pad_sequence(spects, batch_first=True)
+    lengths = torch.LongTensor(lengths)
+    return padded_spects, segs, labels, lengths, fnames
+
+
+def spectral_size(wav_len):
+    layers = [(10,5,0), (8,4,0), (4,2,0), (4,2,0), (4,2,0)]
+    for kernel, stride, padding in layers:
+        wav_len = math.floor((wav_len + 2*padding - 1*(kernel-1) - 1)/stride + 1)
+    return wav_len
+
+
+def get_subset(dataset, percent):
+    A_split = int(len(dataset) * percent)
+    B_split = len(dataset) - A_split
+    dataset, _ = torch.utils.data.random_split(dataset, [A_split, B_split])
+    return dataset
+
+
+class WavPhnDataset(Dataset):
+    def __init__(self, path):
+        self.path = path
+        self.data = list(iter_find_files(self.path, "*.wav"))
+
+    def process_file(self, wav_path):
+        
+        phn_path = wav_path[:-4] + ".phn"
+
+        # load audio
+        audio, sr = torchaudio.load(wav_path)
+        audio = audio[0]
+        audio_len = len(audio)
+        spectral_len = spectral_size(audio_len)
+        len_ratio = (audio_len / spectral_len)
+
+        # load labels -- segmentation and phonemes
+        with open(phn_path, "r") as f:
+            lines = f.readlines()
+            lines = list(map(lambda line: line.split(" "), lines))
+
+            # get segment times
+            times = torch.FloatTensor(list(map(lambda line: int(int(line[1]) / len_ratio), lines)))[:-1]  # don't count end time as boundary
+
+            # get phonemes in each segment (for K times there should be K+1 phonemes)
+            phonemes = list(map(lambda line: line[2].strip(), lines))
+
+        return audio, times.tolist(), phonemes, wav_path
+    
+    def __getitem__(self, idx):
+        audio, seg, phonemes, fname = self.process_file(self.data[idx])
+        audio_len = len(audio)
+        spectral_len = spectral_size(audio_len)
+        len_ratio = (audio_len / spectral_len)
+        return audio, seg, phonemes, audio_len/len_ratio, fname
+
+    def __len__(self):
+        return len(self.data) 
+    
+class TrainTestDataset(WavPhnDataset):
+    @staticmethod
+    def get_datasets(path, val_ratio=0.1, overlap=False, seed: int = 42):
+        """
+        If overlap==False (default) split train into disjoint train/val (random_split).
+        If overlap==True  create val as a Subset sampled from the train dataset
+                           but keep train_dataset as the full set (so val files are also seen in training).
+        """
+        train_full = TrainTestDataset(join(path, 'train'))
+        test_dataset = TrainTestDataset(join(path, 'test'))
+        train_len = len(train_full)
+
+        val_size = int(train_len * val_ratio)
+        if val_size <= 0:
+            # no validation
+            return train_full, None, test_dataset
+
+        if overlap:
+            rng = random.Random(seed)
+            val_indices = rng.sample(range(train_len), val_size)
+            val_dataset = torch.utils.data.Subset(train_full, val_indices)
+            train_dataset = train_full  # full training set (contains val files)
+        else:
+            # exclusive split (current behavior)
+            gen = torch.Generator()
+            gen.manual_seed(seed)
+            train_split = train_len - val_size
+            train_dataset, val_dataset = torch.utils.data.random_split(train_full, [train_split, val_size], generator=gen)
+            # keep .path attribute for compatibility
+            train_dataset.path = join(path, 'train')
+            val_dataset.path = join(path, 'train')
+            return train_dataset, val_dataset, test_dataset
+
+        # ensure compatibility of .path attribute
+        train_dataset.path = join(path, 'train')
+        val_dataset.path = join(path, 'train')
+        return train_dataset, val_dataset, test_dataset
+
+
+class TrainValTestDataset(WavPhnDataset):
+    @staticmethod
+    def get_datasets(path, percent=1.0):
+        train_dataset = TrainValTestDataset(join(path, 'train'))
+        if percent != 1.0:
+            train_dataset = get_subset(train_dataset, percent)
+            train_dataset.path = join(path, 'train')
+        val_dataset   = TrainValTestDataset(join(path, 'val'))
+        test_dataset  = TrainValTestDataset(join(path, 'test'))
+
+        return train_dataset, val_dataset, test_dataset

dutch_preprocess.py

@@ -0,0 +1,197 @@
+import re
+import panphon
+import panphon.distance
+
+ft = panphon.FeatureTable()
+dst = panphon.distance.Distance()
+
+# IFA_TO_IPA = TODO - FUNCTION THAT COMPUTES SCORE
+# NOT MAPPING BUT NEEDS TO BE A PUNCTION
+
+# IPA_TO_LEEHON39 = TODO - NOT MAPPING NEEDS TO BE A FUNCTION BUT ONE TIME - MAYBE IT IS MAPPING?
+
+IFA_TO_IPA = {
+    "p":"p", "b":"b", "t":"t", "d":"d", "k":"k", "g":"ɡ",
+    "f":"f", "v":"v", "s":"s", "z":"z", "h":"h", "x":"x", "G":"ɣ",
+    "m":"m", "n":"n", "N":"ŋ", "l":"l", "r":"r", "w":"ʋ", "j":"j",
+    "S":"ʃ", "Z":"ʒ", "J":"ɲ", "L":"ʎ",
+    "i":"i", "I":"ɪ", "e":"eː", "E":"ɛ", "a":"aː", "A":"ɑ",
+    "o":"oː", "O":"ɔ", "u":"u", "y":"y", "Y":"ʏ", "2":"øː",
+    "9":"œ", "@":"ə", "!" : "ɛi", "V" : "ʌu", "W" : "œy", "h#" : "h#"
+}
+
+LH39_IPA = {
+    "AA": "ɑ", "AE": "æ", "AH": "ʌ", "AO": "ɔ", "AW": "aʊ", "AY": "aɪ",
+    "EH": "ɛ", "ER": "ɝ", "EY": "eɪ", "IH": "ɪ", "IY": "i", "OW": "oʊ",
+    "OY": "ɔɪ", "UH": "ʊ", "UW": "u", "B": "b", "CH": "tʃ", "D": "d",
+    "DH": "ð", "F": "f", "G": "ɡ", "HH": "h", "JH": "dʒ", "K": "k",
+    "L": "l", "M": "m", "N": "n", "NG": "ŋ", "P": "p", "R": "ɹ",
+    "S": "s", "SH": "ʃ", "T": "t", "TH": "θ", "V": "v", "W": "w",
+    "Y": "j", "Z": "z", "ZH": "ʒ"
+}
+
+timit_leehon_39_phonemes = [
+    'ao', 'ae', 'ah','aw', 'er', 'ay', 
+    'b', 'sil', 'ch', 'd', 'dh', 'dx', 'eh', 'el', 'm', 'en', 'ng', 'ey',
+    'f', 'g', 'hh', 'ih', 'iy', 'jh', 'k', 'v', 'w', 'y', 'z', 'sh', 't', 'r', 's', 'th','uh', 'uw', 'oy', 'ow','p'
+]
+
+def get_ipa_from_ifa(ifa_label):
+    
+    if ifa_label.lower() in timit_leehon_39_phonemes:
+        return [ifa_label.lower()]
+    if ifa_label in ['h#', 'tcl']:
+    # if ifa_label in ['h#']:
+        return ["sil"]
+    # if ifa_label in ['tcl']:
+    #     return []
+    
+    # Convert underscores and hyphens to spaces, and remove colons (length is handled by the base vowel mapping or discarded)
+    cleaned = ifa_label.replace(':', '').replace('_', ' ').replace('-', ' ') #.replace('tcl', ' ')
+    # Remove Stress ("), Secondary Stress ('), Syllable dots (.), and nasal tildes (~)
+    cleaned = re.sub(r'[".\'~]', '', cleaned)
+    # parts = cleaned.split()   
+    # parts = cleaned.strip()   
+    parts = cleaned.strip().split()
+    if not parts:
+        return [] 
+    if len(parts) == 1 and len(ifa_label) >1 and ifa_label not in IFA_TO_IPA: #it's a long phoneme label that needs to be splitted to several IPA symbols
+        parts = list(ifa_label)
+    
+    ipa_list = [IFA_TO_IPA.get(p,p) for p in parts if p.strip()]
+    return ipa_list
+
+_leehon39_cache = {}
+
+def find_best_leehon39(target_ipa):
+    # Deterministic IPA->LH39 mapping; memoize since the panphon feature-edit
+    # distance over a fixed inventory recomputes the same answer per occurrence.
+    if target_ipa in _leehon39_cache:
+        return _leehon39_cache[target_ipa]
+    result = _find_best_leehon39(target_ipa)
+    _leehon39_cache[target_ipa] = result
+    return result
+
+def _find_best_leehon39(target_ipa):
+
+    if not target_ipa or target_ipa.strip() == "":
+        return "sil", 0.0
+    
+    if target_ipa.lower() in timit_leehon_39_phonemes:
+        return target_ipa.lower(), 0.0
+    
+    # if target_ipa.lower() in ['h#', 'tcl', 'sil']:
+    if target_ipa.lower() in ['h#', 'sil']:
+        return "sil", 0.0
+    if target_ipa.lower() in ["r", "ɾ"]:
+        return "r", 0.0
+    
+    best_label = "sil"
+    min_dist = 100.0
+    
+    for lh_label, lh_ipa in LH39_IPA.items():
+        
+        d = dst.feature_edit_distance(target_ipa, lh_ipa)
+        if d< min_dist:
+            min_dist = d
+            best_label = lh_label.lower()
+    return best_label, round(min_dist,3)
+
+def aligner_pipeline(ifa_input):
+    ifa_segments = get_ipa_from_ifa(ifa_input)
+    results = []
+    
+    for ipa_seg in ifa_segments:
+        match, d = find_best_leehon39(ipa_seg)
+        results.append( {"ifa_ipa_part" :ipa_seg, "lh39" :match, "dist" :d} )
+    return results
+
+
+import os
+
+# def convert_all_lab_files(directory):
+#     for filename in os.listdir(directory):
+#         if filename.endswith(".lab"):
+#             path = os.path.join(directory, filename)
+#             with open(path, 'r') as f:
+#                 content = f.read().strip()
+            
+#             # Use your existing pipeline logic
+#             # Note: We split the content by space to process each phone
+#             ifa_phones = content.split()
+#             ipa_output = []
+#             for p in ifa_phones:
+#                 # Get the IPA parts from your existing function
+#                 ipa_parts = get_ipa_from_ifa(p)
+#                 ipa_output.extend(ipa_parts)
+            
+#             # Join with spaces and write back
+#             new_content = " ".join(ipa_output)
+#             with open(path, 'w') as f:
+#                 f.write(new_content)
+#     print(f"Done! All .lab files in {directory} converted to IPA.")
+
+# # Run this in your main block
+# # convert_all_lab_files('/home/rotem/projects/datasets/IFA_dutch_split/test')
+
+# # convert_all_lab_files('/home/rotem/projects/datasets/IFA_dutch_split/test')
+
+
+
+import os
+
+def create_lab_files(phn_folder, lab_folder):
+    if not os.path.exists(lab_folder):
+        os.makedirs(lab_folder)
+
+    for filename in os.listdir(phn_folder):
+        if filename.endswith(".phn"):
+            with open(os.path.join(phn_folder, filename), 'r') as f:
+                lines = f.readlines()
+            
+            ipa_sequence = []
+            for line in lines:
+                parts = line.strip().split()
+                if len(parts) < 3: continue
+                
+                label = parts[2]
+                # Use your existing mapping function
+                ipa_symbols = get_ipa_from_ifa(label) 
+                
+                # Filter out 'sil' if you want MFA to handle silence automatically, 
+                # but usually keeping them is fine for phone-level alignment.
+                ipa_sequence.extend(ipa_symbols)
+            
+            # Save to .lab file (space separated string)
+            lab_filename = filename.replace(".phn", ".lab")
+            with open(os.path.join(lab_folder, lab_filename), 'w') as f:
+                f.write(" ".join(ipa_sequence))
+                
+def generate_ipa_lexicon(all_ipa_symbols, output_path):
+    with open(output_path, 'w') as f:
+        # Add a silence mapping just in case
+        f.write("sil\tsil\n")
+        # Map every unique IPA symbol to itself
+        for symbol in sorted(list(set(all_ipa_symbols))):
+            if symbol != "sil":
+                f.write(f"{symbol}\t{symbol}\n")
+
+# Run it
+# create_lab_files("/home/rotem/projects/datasets/IFA_dutch_split/test", "/home/rotem/projects/datasets/IFA_dutch_split/test")
+
+if __name__ == "__main__":
+    test_cases = ["sil n Y l sil e: n sil t w e: sil d r i sil v i r sil v Ei f sil z E s sil z e: v @ n sil A x t sil n e: x @ sil t i n sil E l f sil t w a: l f sil n Y l sil sil"]
+    # test_cases = ["x@l", "@-r-h-a", "e:-j", "r9y", "ao", "sil", "@", "E", "he:l-@_hAr", "t_b", "o:", "N"]
+    for case in test_cases:
+        print(f"\nINPUT: {case}")
+        output = aligner_pipeline(case)
+        # [x["lh39"] for x in output]
+        if not output:
+            print("Results: None")
+        else:
+            for item in output:
+                print(f" Mapped '{item['ifa_ipa_part']}' -> {item['lh39']} (dist_score: {item['dist']})")
+    # convert_all_lab_files('/home/rotem/projects/datasets/IFA_dutch_split/test')
+    # Run it
+    create_lab_files("/home/rotem/projects/datasets/IFA_dutch_split/test", "/home/rotem/projects/datasets/IFA_dutch_split/test")
+                

falcon_viz.py

@@ -0,0 +1,417 @@
+"""
+falcon_viz.py — Time-aligned alignment visualizations for FALCON.
+
+Two entry points, both sharing the same per-panel drawers so the styling stays
+identical:
+
+    make_alignment_panels(wav, ckpt, out_path, ...)
+        -> ONE tall figure, 5 panels stacked on a shared time axis (for the
+           README / a downloadable overview).
+
+    make_alignment_separate(wav, ckpt, out_dir, ...)
+        -> a LIST of (path, caption), one wide, full-size figure per panel (for
+           the web app, where a single stacked figure renders too small to read).
+
+Panels: waveform, log-mel spectrogram, phoneme posteriors, Soft-DP cost matrix +
+alignment path, and the contrastive boundary score with its detected peaks. The
+same predicted boundaries (crimson) and truth boundaries (charcoal) are drawn on
+every panel so the time alignment is visually verifiable.
+
+This module does NOT modify any model/training file. It re-runs the model forward
+(replicating predict.py:main_predict) and reproduces utils.phoneme_alignment's DP
+matrix locally (re-using utils.compute_phi_1 / compute_phi_2 unchanged).
+"""
+
+import os
+from glob import glob
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torchaudio
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+from matplotlib.gridspec import GridSpec
+
+from dataloader import spectral_size
+from predict import _load_model
+from utils import (
+    max_min_norm,
+    get_timit_61_phoneme_mappings,
+    compute_phi_1,
+    compute_phi_2,
+    phoneme_to_idx_MACRO,
+    timit_to_leehon_map_MACRO,
+    timit_leehon_39_phonemes,
+)
+
+# Hard-coded in the model forward (next_frame_classifier.py) — CNN total stride.
+# One latent frame ~= len_ratio audio samples ~= 10 ms at 16 kHz.
+LEN_RATIO = 161.34011627906978
+SR = 16000
+
+# Light / pastel palette so the crimson boundary overlays stand out.
+PRED_COLOR = "crimson"
+TRUTH_COLOR = "#37474f"        # charcoal (distinct from the blue/purple cmaps)
+WAVE_COLOR = "#3b6fb6"         # blue (matches the example waveform figure)
+SCORE_COLOR = "#3a7ca5"
+SPEC_CMAP = "viridis"
+POST_CMAP = "viridis"
+DP_CMAP = "viridis"
+
+
+# --------------------------------------------------------------------------- #
+#  DP matrix reconstruction (mirrors utils.phoneme_alignment, returns matrix)  #
+# --------------------------------------------------------------------------- #
+def _build_dp_matrix(p_seq, w_phi, original_lengths, derivative_preds_np, probs_real):
+    """Re-run the exact Soft-DP forward + soft-argmax backtrack from
+    utils.phoneme_alignment, additionally returning the DP matrix so it can be
+    plotted (the stock function only returns the path)."""
+    gamma = 1e-20
+    T = int(original_lengths[0])
+    n = len(p_seq)
+    device = derivative_preds_np.device
+
+    if isinstance(probs_real, np.ndarray):
+        probs_real = torch.tensor(probs_real, device=device)
+    cumsum_probs = torch.cumsum(probs_real, dim=0)
+
+    phoneme_mappings = {
+        p.lower(): timit_to_leehon_map_MACRO.get(p.lower(), "sil")
+        if p.lower() not in timit_leehon_39_phonemes
+        else p.lower()
+        for p in p_seq
+    }
+    derivatives = torch.cat(
+        [torch.tensor([0], device=device), torch.diff(derivative_preds_np, dim=0)]
+    )
+
+    dp_mat = torch.full((n, T, T), float(-1e9), device=device)
+
+    t_e = torch.arange(T, device=device)
+    dp_mat[0, 0, :] = (
+        w_phi[0] * compute_phi_1(derivatives, 0, t_e)
+        + w_phi[1] * compute_phi_1(derivatives, 0, t_e)
+    )
+
+    for i in range(1, n):
+        p_idx = phoneme_to_idx_MACRO[phoneme_mappings[p_seq[i].lower()]]
+        t_start = torch.arange(T, device=device)
+        t_end = torch.arange(T, device=device)
+        t_start_grid, t_end_grid = torch.meshgrid(t_start, t_end, indexing="ij")
+        valid_mask = t_start_grid < t_end_grid
+
+        phi1_dev = compute_phi_1(derivatives, t_start_grid, t_end_grid)
+        phi2 = compute_phi_2(cumsum_probs, p_idx, t_start_grid, t_end_grid)
+        total_phi = w_phi[0] * phi1_dev + w_phi[1] * phi2
+
+        col_lse = torch.logsumexp(dp_mat[i - 1] / gamma, dim=0) * gamma
+        prev_scores = torch.where(
+            valid_mask,
+            col_lse.unsqueeze(1).expand(T, T),
+            torch.full((T, T), float(-1e9), device=device),
+        )
+        dp_mat[i] = torch.where(
+            valid_mask, total_phi + prev_scores, torch.full_like(total_phi, float(-1e9))
+        )
+
+    best_start_times = torch.zeros((n), dtype=derivative_preds_np.dtype, device=device)
+    best_prev_t_end = T - 1
+    for i in range(n):
+        cur_ph = n - 1 - i
+        scores = dp_mat[cur_ph, :, best_prev_t_end]
+        soft_weights = torch.softmax(scores / gamma, dim=0)
+        expected_idx = (
+            soft_weights
+            * torch.arange(T, device=device, dtype=derivative_preds_np.dtype)
+        ).sum()
+        best_start_times[cur_ph] = expected_idx
+        best_prev_t_end = int(expected_idx.round().item())
+
+    dp_to_plot = dp_mat.detach().cpu().max(dim=1)[0].numpy()  # (n, T)
+    best_start_frames = best_start_times.detach().cpu().numpy()
+    return dp_to_plot, best_start_frames
+
+
+# --------------------------------------------------------------------------- #
+#  Forward pass + array extraction (replicates predict.py:main_predict)         #
+# --------------------------------------------------------------------------- #
+def _extract_arrays(wav, ckpt, annotation):
+    model, _peak_params = _load_model(ckpt)
+    model.eval()
+
+    audio, sr = torchaudio.load(wav)
+    assert sr == SR, "model expects 16 kHz audio"
+    audio = audio[0]
+    audio_len = len(audio)
+    spectral_len = spectral_size(audio_len)
+    len_ratio = audio_len / spectral_len  # ~= LEN_RATIO
+
+    base_dir = os.path.dirname(wav)
+    base_name = os.path.basename(wav).split(".")[0]
+    matches = glob(os.path.join(base_dir, f"{base_name}*.{annotation}"))
+    phn_path = matches[0] if matches else wav.replace("wav", "phn")
+
+    with open(phn_path, "r") as f:
+        lines = [ln.split() for ln in f.readlines()]
+    truth_secs = [float(ln[1]) / SR for ln in lines][:-1]
+    phonemes = [ln[2].strip() for ln in lines]
+    truth_labels = list(phonemes)
+
+    length = [audio_len / len_ratio]
+    with torch.no_grad():
+        preds, original_lengths, probs, frame_labels, seg, total_peaks, w_phi = model(
+            audio.unsqueeze(0), None, [phonemes], length
+        )
+
+    # Contrastive / latent boundary score (predict.py 168-172).
+    p = preds[1][0]
+    p = max_min_norm(p)
+    p_np = p.detach().numpy()[0]
+    p_np = p_np - np.median(p_np)
+
+    # Phoneme posteriors (predict.py 143).
+    probs_real = F.softmax(probs, dim=-1).squeeze(0).detach().numpy()  # (T, 39)
+    _, idx_to_phoneme = get_timit_61_phoneme_mappings()
+    phoneme_labels = [idx_to_phoneme[i] for i in range(39)]
+
+    pred_secs = list(total_peaks[0])
+
+    # Phoneme labels at predicted-segment midpoints (segment i = [bound[i],
+    # bound[i+1]] with bounds = predicted boundaries + the utterance end), matching
+    # how the TextGrid assigns each phoneme to a predicted interval.
+    bounds = [float(x) for x in pred_secs] + [audio_len / SR]
+    seg_mids = [(bounds[i] + bounds[i + 1]) / 2.0 for i in range(len(bounds) - 1)]
+    n_lab = min(len(seg_mids), len(phonemes))
+    label_mids = seg_mids[:n_lab]
+    label_text = phonemes[:n_lab]
+
+    # DP matrix (reproduce phoneme_alignment locally).
+    w_phi_vec = w_phi.detach()
+    deriv_arg = torch.tensor(p_np, dtype=torch.float32)
+    dp_to_plot, dp_path_frames = _build_dp_matrix(
+        phonemes, w_phi_vec, [int(original_lengths[0])], deriv_arg, probs_real
+    )
+
+    return dict(
+        audio=audio.numpy(),
+        sr=sr,
+        len_ratio=len_ratio,
+        duration=audio_len / SR,
+        latent_score=p_np,
+        probs_real=probs_real,
+        phoneme_labels=phoneme_labels,
+        pred_secs=pred_secs,
+        truth_secs=truth_secs,
+        truth_labels=truth_labels,
+        dp_to_plot=dp_to_plot,
+        dp_path_frames=dp_path_frames,
+        phonemes=phonemes,
+        label_mids=label_mids,
+        label_text=label_text,
+    )
+
+
+# --------------------------------------------------------------------------- #
+#  Shared per-panel drawers                                                     #
+# --------------------------------------------------------------------------- #
+def _overlay_boundaries(ax, pred_secs, truth_secs, show_truth, label_first=False):
+    """Predicted (crimson dashed) + optional truth (charcoal dotted) lines."""
+    for j, t in enumerate(truth_secs if show_truth else []):
+        ax.axvline(t, color=TRUTH_COLOR, linestyle=":", linewidth=0.9, alpha=0.6,
+                   label="Truth boundary" if (label_first and j == 0) else None, zorder=2)
+    for j, t in enumerate(pred_secs):
+        ax.axvline(t, color=PRED_COLOR, linestyle="--", linewidth=1.1, alpha=0.9,
+                   label="Predicted boundary" if (label_first and j == 0) else None, zorder=3)
+
+
+def _annotate_phoneme_tier(ax, d):
+    """Write the input phoneme labels at their predicted-segment midpoints, just
+    below the x-axis (a phoneme tier)."""
+    mids = d.get("label_mids") or []
+    labels = d.get("label_text") or []
+    trans = ax.get_xaxis_transform()  # x in data coords, y in axes fraction
+    for m, lab in zip(mids, labels):
+        ax.text(m, -0.07, str(lab), transform=trans, rotation=90, ha="center",
+                va="top", fontsize=5.5, color="#333333", clip_on=False)
+
+
+def _panel_waveform(ax, d, show_truth, label_first=True):
+    audio = d["audio"]
+    dur = d["duration"]
+    t = np.linspace(0, dur, num=len(audio))
+    ax.plot(t, audio, color=WAVE_COLOR, linewidth=0.5)
+    ax.set_ylabel("Amplitude")
+    ax.margins(x=0)
+    ymax = (float(np.abs(audio).max()) or 1.0) * 1.15
+    ax.set_ylim(-ymax, ymax)
+    _overlay_boundaries(ax, d["pred_secs"], d["truth_secs"], show_truth, label_first=label_first)
+    if label_first:
+        ax.legend(loc="upper right", fontsize=8, framealpha=0.9, ncol=2)
+
+
+def _panel_spectrogram(ax, d, show_truth):
+    audio = d["audio"]
+    dur = d["duration"]
+    mel = torchaudio.transforms.MelSpectrogram(
+        sample_rate=SR, n_fft=400, hop_length=160, n_mels=80
+    )(torch.tensor(audio).float().unsqueeze(0))
+    mel_db = torchaudio.transforms.AmplitudeToDB(top_db=80)(mel).squeeze(0).numpy()
+    ax.imshow(mel_db, aspect="auto", origin="lower",
+              extent=[0, dur, 0, SR / 2 / 1000.0], cmap=SPEC_CMAP)
+    ax.set_ylabel("Freq (kHz)")
+    _overlay_boundaries(ax, d["pred_secs"], d["truth_secs"], show_truth)
+
+
+def _panel_posteriors(ax, d, show_truth, colorbar=True):
+    probs_real = d["probs_real"]
+    dur = d["duration"]
+    im = ax.imshow(probs_real.T, aspect="auto", origin="lower",
+                   extent=[0, dur, -0.5, 38.5], cmap=POST_CMAP, interpolation="nearest")
+    ax.set_yticks(range(39))
+    ax.set_yticklabels(d["phoneme_labels"], fontsize=5.5)
+    ax.set_ylabel("Phoneme (LH-39)")
+    if colorbar:
+        ax.figure.colorbar(im, ax=ax, label="P(phoneme)", pad=0.01, fraction=0.025)
+    _overlay_boundaries(ax, d["pred_secs"], d["truth_secs"], show_truth)
+
+
+def _panel_dp(ax, d, show_truth, colorbar=True):
+    dp_to_plot = d["dp_to_plot"]
+    dur = d["duration"]
+    len_ratio = d["len_ratio"]
+    masked = np.ma.masked_where(dp_to_plot <= -1e8, dp_to_plot)
+    cmap = getattr(plt.cm, DP_CMAP).copy()
+    cmap.set_bad(color="white")
+    n_ph = dp_to_plot.shape[0]
+    im = ax.imshow(masked, aspect="auto", origin="lower",
+                   extent=[0, dur, -0.5, n_ph - 0.5], cmap=cmap, interpolation="nearest")
+    path_secs = np.asarray(d["dp_path_frames"]) * len_ratio / SR
+    ax.plot(path_secs, np.arange(n_ph), color=PRED_COLOR, marker=".", markersize=4,
+            linewidth=1.2, label="Optimal alignment path")
+    ax.set_ylabel("Phoneme position")
+    if colorbar:
+        ax.figure.colorbar(im, ax=ax, label="DP score", pad=0.01, fraction=0.025)
+    _overlay_boundaries(ax, d["pred_secs"], d["truth_secs"], show_truth)
+    ax.legend(loc="lower right", fontsize=8, framealpha=0.9)
+
+
+def _panel_contrastive(ax, d, show_truth):
+    # Boundary score (red) + its derivative (magenta), as in predict.py's run plot,
+    # but (1) auto-scaled robustly so the per-boundary structure is visible instead
+    # of being squashed by the silence->speech onset spike, and (2) without the
+    # redundant predicted-boundary markers — the predicted boundaries are the red
+    # dashed lines shared across every panel. Ground truth omitted (not available
+    # at inference). x-axis is time (s).
+    s = np.asarray(d["latent_score"], dtype=float)
+    len_ratio = d["len_ratio"]
+    n = len(s)
+    t = np.arange(n) * len_ratio / SR
+    deriv = np.concatenate([[0.0], np.diff(s)])
+
+    ax.plot(t, deriv, marker="o", markersize=1.6, linewidth=0.7, alpha=0.8,
+            color="magenta", label="Derivative of latent score")
+    ax.plot(t, s, marker="*", markersize=2.2, linewidth=0.8, color="red",
+            label="Latent score")
+
+    # Robust symmetric y-limit: zoom to the 96th-percentile magnitude so the small
+    # per-boundary structure fills the panel; the rare large onset spike clips off.
+    mag = np.concatenate([np.abs(s), np.abs(deriv)])
+    A = max(float(np.percentile(mag, 96)) * 1.5, 0.03) if mag.size else 0.1
+    ax.set_ylim(-A, A)
+    ax.set_ylabel("Score")
+    ax.margins(x=0)
+    _overlay_boundaries(ax, d["pred_secs"], [], show_truth=False)
+    ax.legend(loc="upper right", fontsize=7, framealpha=0.9, ncol=2)
+
+
+# Order shared by both makers: (key, caption, drawer, separate-figure height).
+_PANELS = [
+    ("waveform",    "1. Waveform",                                  _panel_waveform,     2.6),
+    ("spectrogram", "2. Log-mel spectrogram",                       _panel_spectrogram,  2.8),
+    ("posteriors",  "3. Phoneme posteriors",                        _panel_posteriors,   3.6),
+    ("dp",          "4. Soft-DP cost matrix + alignment path",      _panel_dp,           3.4),
+    ("contrastive", "5. Contrastive boundary score", _panel_contrastive, 2.6),
+]
+
+
+# --------------------------------------------------------------------------- #
+#  Combined figure (README)                                                    #
+# --------------------------------------------------------------------------- #
+def make_alignment_panels(wav, ckpt, out_path, w_phi=0.5, language="english",
+                          annotation="phn", show_truth=True):
+    """Build the stacked, time-aligned multi-panel figure and save it to out_path."""
+    if language != "english":
+        print(f"[falcon_viz] language='{language}' not supported; using english path.")
+    d = _extract_arrays(wav, ckpt, annotation)
+    dur = d["duration"]
+
+    fig = plt.figure(figsize=(12, 17), dpi=150, constrained_layout=True)
+    gs = GridSpec(5, 1, figure=fig, height_ratios=[1.0, 1.3, 1.7, 1.6, 1.0])
+    ax_wave = fig.add_subplot(gs[0])
+    ax_spec = fig.add_subplot(gs[1], sharex=ax_wave)
+    ax_post = fig.add_subplot(gs[2], sharex=ax_wave)
+    ax_dp = fig.add_subplot(gs[3], sharex=ax_wave)
+    ax_score = fig.add_subplot(gs[4], sharex=ax_wave)
+    axes = [ax_wave, ax_spec, ax_post, ax_dp, ax_score]
+
+    _panel_waveform(ax_wave, d, show_truth, label_first=True)
+    _panel_spectrogram(ax_spec, d, show_truth)
+    _panel_posteriors(ax_post, d, show_truth, colorbar=True)
+    _panel_dp(ax_dp, d, show_truth, colorbar=True)
+    _panel_contrastive(ax_score, d, show_truth)
+
+    for ax, (_key, caption, _drawer, _h) in zip(axes, _PANELS):
+        ax.set_title(caption, loc="left", fontweight="bold", fontsize=11)
+        ax.tick_params(labelbottom=True)
+        _annotate_phoneme_tier(ax, d)
+    ax_score.set_xlabel("Time (s)", fontsize=12)
+    ax_wave.set_xlim(0, dur)
+
+    fig.suptitle("FALCON forced-alignment — time-aligned representations",
+                 fontsize=14, fontweight="bold")
+    fig.savefig(out_path, dpi=150, bbox_inches="tight")
+    plt.close(fig)
+    print(f"[falcon_viz] saved {out_path}")
+    return out_path
+
+
+# --------------------------------------------------------------------------- #
+#  Separate per-panel figures (web app)                                        #
+# --------------------------------------------------------------------------- #
+def make_alignment_separate(wav, ckpt, out_dir, w_phi=0.5, language="english",
+                            annotation="phn", show_truth=True):
+    """Build one wide, full-size figure per panel; return [(path, caption), ...]."""
+    if language != "english":
+        print(f"[falcon_viz] language='{language}' not supported; using english path.")
+    d = _extract_arrays(wav, ckpt, annotation)
+    dur = d["duration"]
+    os.makedirs(out_dir, exist_ok=True)
+
+    out = []
+    for key, caption, drawer, h in _PANELS:
+        fig, ax = plt.subplots(figsize=(12, h), dpi=130)
+        if key == "waveform":
+            drawer(ax, d, show_truth, label_first=True)
+        else:
+            drawer(ax, d, show_truth)
+        ax.set_title(caption, loc="left", fontweight="bold", fontsize=12)
+        ax.set_xlim(0, dur)
+        _annotate_phoneme_tier(ax, d)
+        ax.set_xlabel("Time (s)", labelpad=26)
+        p = os.path.join(out_dir, f"panel_{key}.png")
+        fig.savefig(p, dpi=130, bbox_inches="tight")
+        plt.close(fig)
+        out.append((p, caption))
+    return out
+
+
+if __name__ == "__main__":
+    _here = os.path.dirname(os.path.abspath(__file__))
+    make_alignment_panels(
+        wav=os.path.join(_here, "assets", "fasw0sa2.wav"),
+        ckpt=os.path.join(_here, "pretrained_models", "falcon_timit_english.pt"),
+        out_path=os.path.join(_here, "assets", "example_panels.png"),
+    )

mfa_assets/extracted_models/g2p/english_us_arpa_g2p/phones.sym

@@ -0,0 +1,89 @@
+<eps> 0
+sil 1
+spn 2
+AA0 3
+AA1 4
+AA2 5
+AE0 6
+AE1 7
+AE2 8
+AH0 9
+AH1 10
+AH2 11
+AO0 12
+AO1 13
+AO2 14
+AW0 15
+AW1 16
+AW2 17
+AY0 18
+AY1 19
+AY2 20
+B 21
+CH 22
+D 23
+DH 24
+EH0 25
+EH1 26
+EH2 27
+ER0 28
+ER1 29
+ER2 30
+EY0 31
+EY1 32
+EY2 33
+F 34
+G 35
+HH 36
+IH0 37
+IH1 38
+IH2 39
+IY0 40
+IY1 41
+IY2 42
+JH 43
+K 44
+L 45
+M 46
+N 47
+NG 48
+OW0 49
+OW1 50
+OW2 51
+OY0 52
+OY1 53
+OY2 54
+P 55
+R 56
+S 57
+SH 58
+T 59
+TH 60
+UH0 61
+UH1 62
+UH2 63
+UW0 64
+UW1 65
+UW2 66
+V 67
+W 68
+Y 69
+Z 70
+ZH 71
+#0 72
+#1 73
+#2 74
+#3 75
+#4 76
+#5 77
+#6 78
+#7 79
+#8 80
+#9 81
+#10 82
+#11 83
+#12 84
+#13 85
+#14 86
+#15 87
+#16 88

mfa_g2p.py

@@ -0,0 +1,199 @@
+"""
+MFA-compatible word -> Lee-Hon-39 phoneme front-end for *word-level* alignment.
+
+This is the apples-to-apples counterpart of `word_g2p.py` (which uses espeak):
+instead of espeak it phonemizes orthographic words with the **same** open-source
+G2P models that Montreal Forced Aligner uses, then maps the result into LH39:
+
+    word --MFA pronunciation dictionary lookup (highest-prob entry)--> phones
+    word (OOV, English only) --english_us_arpa pynini G2P WFST--> ARPAbet
+    English ARPAbet --strip stress, lowercase, timit_to_leehon_map_MACRO--> LH39
+    Dutch/German IPA --panphon articulatory distance (dutch_preprocess)--> LH39
+
+MFA at alignment time looks each word up in its dictionary first and only invokes
+the G2P WFST for out-of-vocabulary words; this module mirrors that. Selecting this
+backend (vs espeak) isolates the G2P front-end as the only thing that differs from
+MFA, so a word-level FDNFA-vs-MFA comparison measures the *aligner*, not the
+phonemizer.
+
+Language is chosen with FDNFA_G2P_VOICE (the same env word_g2p uses):
+  en-us -> english_us_arpa (ARPAbet, pynini OOV)   nl -> dutch_cv (IPA)
+  de    -> german_mfa (IPA)                         he -> no MFA model (espeak fallback)
+"""
+import os
+import subprocess
+
+from utils import timit_to_leehon_map_MACRO, timit_leehon_39_phonemes
+
+# Default MFA model locations (the standard `mfa model download` cache).
+_MFA_ROOT = os.environ.get("MFA_ROOT_DIR", os.path.expanduser("~/Documents/MFA"))
+_DICT_DIR = os.path.join(_MFA_ROOT, "pretrained_models", "dictionary")
+_ARPA_G2P_DIR = os.path.join(_MFA_ROOT, "extracted_models", "g2p", "english_us_arpa_g2p")
+ARPA_G2P_FST = os.environ.get("FDNFA_MFA_G2P_FST", os.path.join(_ARPA_G2P_DIR, "model.fst"))
+ARPA_G2P_PHONES = os.environ.get("FDNFA_MFA_G2P_PHONES", os.path.join(_ARPA_G2P_DIR, "phones.sym"))
+# conda env that has pynini installed (for OOV G2P on the WFST).
+MFA_ENV_PY = os.environ.get("FDNFA_MFA_ENV_PY", os.path.expanduser("~/miniconda3/envs/aligner/bin/python"))
+
+# voice -> (dictionary file, phone alphabet). FDNFA_MFA_DICT overrides the dict.
+_LANG_CFG = {
+    "en-us": ("english_us_arpa.dict", "arpa"),
+    "en":    ("english_us_arpa.dict", "arpa"),
+    "nl":    ("dutch_cv.dict",        "ipa"),
+    "de":    ("german_mfa.dict",      "ipa"),
+}
+_VOICE = os.environ.get("FDNFA_G2P_VOICE", "en-us")
+
+
+def _cfg_for_voice(voice):
+    """(dictionary path, phone alphabet) for an espeak-style voice code.
+    FDNFA_MFA_DICT overrides the dictionary path for the default voice only."""
+    voice = voice or _VOICE
+    dict_file, alphabet = _LANG_CFG.get(voice, ("english_us_arpa.dict", "arpa"))
+    if voice == _VOICE and os.environ.get("FDNFA_MFA_DICT"):
+        return os.environ["FDNFA_MFA_DICT"], alphabet
+    return os.path.join(_DICT_DIR, dict_file), alphabet
+
+
+def mfa_available(voice=None):
+    """True if an MFA pronunciation dictionary exists locally for this language.
+    Callers (e.g. the app) use this to decide whether to use the MFA-like G2P or
+    fall back to espeak when the dictionary isn't installed."""
+    dict_path, _ = _cfg_for_voice(voice)
+    return os.path.exists(dict_path)
+
+
+# Backwards-compatible module-level defaults (the default voice's config).
+_DICT_FILE, _ALPHABET = _LANG_CFG.get(_VOICE, ("english_us_arpa.dict", "arpa"))
+ARPA_DICT, _ = _cfg_for_voice(_VOICE)
+
+# Reuse the exact closure-insertion rule from the espeak front-end so the only
+# thing differing between the espeak and MFA word backends is the G2P.
+from word_g2p import USE_CLOSURES, _with_closures
+
+_dicts = {}           # voice -> {word_lower: [phones]}  (highest-prob entry)
+_cache = {}           # (word_lower, voice) -> [lh39, ...]
+_oov = set()          # words not found in any dictionary (for reporting)
+
+
+def _load_dict(voice=None):
+    """Parse the MFA dictionary for `voice` once (cached per voice). Format:
+    word <tab> [prob cols <tab>] PHONES, where PHONES (final tab-separated field)
+    is space-separated phones and the first float column (when present) is the
+    pronunciation probability.
+
+    Like MFA, keep the **highest-probability** pronunciation per word (MFA's most-
+    likely variant; it then disambiguates acoustically, which we cannot). Entries
+    with no probability column are treated as probability 1.0."""
+    voice = voice or _VOICE
+    if voice in _dicts:
+        return _dicts[voice]
+    dict_path, _alpha = _cfg_for_voice(voice)
+    best = {}   # word -> (prob, phones)
+    with open(dict_path, "r", encoding="utf-8") as f:
+        for line in f:
+            line = line.rstrip("\n")
+            if not line:
+                continue
+            parts = line.split("\t")
+            if len(parts) < 2:
+                continue
+            word = parts[0].lower()
+            phones = parts[-1].split()
+            try:
+                prob = float(parts[1]) if len(parts) >= 3 else 1.0
+            except ValueError:
+                prob = 1.0
+            if word and (word not in best or prob > best[word][0]):
+                best[word] = (prob, phones)
+    _dicts[voice] = {w: ph for w, (_, ph) in best.items()}
+    return _dicts[voice]
+
+
+def arpa_to_lh39(phones):
+    """ARPAbet (with stress digits) -> LH39, dropping non-phone tokens (spn/sil)."""
+    out = []
+    for p in phones:
+        base = p.rstrip("0123456789").lower()   # AH0 -> ah, B -> b
+        if base in ("spn", "sil", "sp", ""):
+            continue
+        if base in timit_leehon_39_phonemes:
+            out.append(base)
+        else:
+            out.append(timit_to_leehon_map_MACRO.get(base, "sil"))
+    return out
+
+
+def ipa_to_lh39(phones):
+    """IPA phones (dutch_cv / german_mfa dicts) -> LH39 via panphon distance — the
+    same articulatory mapping the espeak/phoneme paths use (dutch_preprocess)."""
+    import dutch_preprocess
+    out = []
+    for p in phones:
+        if p in ("spn", "sil", "sp", ""):
+            continue
+        out.append(dutch_preprocess.find_best_leehon39(p)[0])
+    return out
+
+
+def _g2p_oov(word):
+    """Phonemize an OOV English word with the english_us_arpa pynini WFST (run in
+    the mfa env, which has pynini). Mirrors MFA's own G2P: compose the word
+    acceptor with the pair-n-gram model and take the shortest path, decoded via
+    phones.sym. Returns a list of ARPAbet phones, or [] if unavailable."""
+    if not (os.path.exists(MFA_ENV_PY) and os.path.exists(ARPA_G2P_FST)
+            and os.path.exists(ARPA_G2P_PHONES)):
+        return []
+    code = (
+        "import sys,pynini\n"
+        f"fst=pynini.Fst.read({ARPA_G2P_FST!r})\n"
+        f"ps=pynini.SymbolTable.read_text({ARPA_G2P_PHONES!r})\n"
+        "fst.set_output_symbols(ps)\n"
+        "w=sys.argv[1].lower()\n"
+        "try:\n"
+        "    lat=pynini.compose(pynini.accep(w, token_type='utf8'), fst)\n"
+        "    print(pynini.shortestpath(lat).string(ps))\n"
+        "except Exception:\n"
+        "    print('')\n"
+    )
+    try:
+        out = subprocess.run([MFA_ENV_PY, "-c", code, word],
+                             capture_output=True, text=True, timeout=30).stdout
+        return out.strip().split()
+    except Exception:
+        return []
+
+
+def word_to_lh39_mfa(word, voice=None):
+    """Orthographic word -> list of LH39 phonemes via the MFA G2P for `voice`
+    (en-us/en -> english_us_arpa + pynini OOV; de -> german_mfa; nl -> dutch_cv).
+    `voice=None` uses the env default (FDNFA_G2P_VOICE), preserving the original
+    single-language behaviour."""
+    voice = voice or _VOICE
+    _dict_path, alphabet = _cfg_for_voice(voice)
+    key = (word.lower(), voice)
+    if key in _cache:
+        return _cache[key]
+    d = _load_dict(voice)
+    phones = d.get(word.lower())
+    if phones is None:
+        _oov.add(word.lower())
+        # English OOV -> MFA's pynini WFST. German OOV is pre-resolved in the
+        # merged dictionary. Any word still unresolved (e.g. Dutch, which has no
+        # MFA G2P model) becomes a single 'sil', exactly as MFA treats an unknown
+        # word.
+        if alphabet == "arpa":
+            phones = _g2p_oov(word.lower())
+    if alphabet == "arpa":
+        lh39 = arpa_to_lh39(phones) if phones else []
+    else:
+        lh39 = ipa_to_lh39(phones) if phones else []
+    if not lh39:
+        lh39 = ["sil"]
+    if USE_CLOSURES:
+        lh39 = _with_closures(lh39)
+    _cache[key] = lh39
+    return lh39
+
+
+def oov_words():
+    return set(_oov)

next_frame_classifier.py

@@ -0,0 +1,467 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import hydra
+from utils import LambdaLayer, PrintShapeLayer, length_to_mask, get_timit_61_phoneme_mappings, timit_to_leehon
+from dataloader import TrainTestDataset
+from collections import defaultdict
+import random
+from utils import max_min_norm, detect_peaks, create_truth_probs_real
+import torch.nn.utils.rnn as rnn_utils
+from memory_profiler import profile
+
+class NextFrameClassifier(nn.Module):
+    def __init__(self, hp):
+        super(NextFrameClassifier, self).__init__()
+        self.w_phi = nn.Parameter(torch.tensor([0.5, 0.5], dtype=torch.float32))
+        self.w_pos_neg = nn.Parameter(torch.tensor([0.5, 0.5], dtype=torch.float32))
+        self.hp = hp
+
+        Z_DIM = hp.z_dim
+        LS = hp.latent_dim if hp.latent_dim != 0 else Z_DIM
+
+        self.phoneme_to_index, self.idx_to_phoneme = get_timit_61_phoneme_mappings()
+
+        self.enc = nn.Sequential(
+            nn.Conv1d(1, LS, kernel_size=10, stride=5, padding=0, bias=False),
+            nn.BatchNorm1d(LS),
+            nn.LeakyReLU(),
+            nn.Conv1d(LS, LS, kernel_size=8, stride=4, padding=0, bias=False),
+            nn.BatchNorm1d(LS),
+            nn.LeakyReLU(),
+            nn.Conv1d(LS, LS, kernel_size=4, stride=2, padding=0, bias=False),
+            nn.BatchNorm1d(LS),
+            nn.LeakyReLU(),
+            nn.Conv1d(LS, LS, kernel_size=4, stride=2, padding=0, bias=False),
+            nn.BatchNorm1d(LS),
+            nn.LeakyReLU(),
+            nn.Conv1d(LS, Z_DIM, kernel_size=4, stride=2, padding=0, bias=False),
+            LambdaLayer(lambda x: x.transpose(1,2)),
+        )
+        print("learning features from raw wav")
+        
+        if self.hp.z_proj != 0:
+            if self.hp.z_proj_linear:
+                self.enc.add_module(
+                    "z_proj",
+                    nn.Sequential(
+                        nn.Dropout1d(self.hp.z_proj_dropout),
+                        nn.Linear(Z_DIM, self.hp.z_proj),
+                    )
+                )
+            else:
+                self.enc.add_module(
+                    "z_proj",
+                    nn.Sequential(
+                        nn.Dropout1d(self.hp.z_proj_dropout),
+                        nn.Linear(Z_DIM, Z_DIM), nn.LeakyReLU(),
+                        nn.Dropout1d(self.hp.z_proj_dropout),
+                        nn.Linear(Z_DIM, self.hp.z_proj),
+                    )
+                )
+        self.pred_steps = list(range(1 + self.hp.pred_offset, 1 + self.hp.pred_offset + self.hp.pred_steps))
+        print(f"prediction steps: {self.pred_steps}")
+        
+        self.bi_lstm = nn.LSTM(
+            input_size=self.hp.z_proj,
+            hidden_size=512,
+            num_layers=5,#3,
+            bidirectional=True,
+            batch_first=True
+        )
+        def init_weights(m):
+            if isinstance(m, nn.LSTM):
+                for name, param in m.named_parameters():
+                    if 'weight' in name:
+                        nn.init.xavier_uniform_(param.data)
+                    elif 'bias' in name:
+                        nn.init.zeros_(param.data)
+        self.bi_lstm.apply(init_weights)
+        self.fc = nn.Linear(512 * 2, hp.num_classes)
+
+    def score(self, f, b):
+        return F.cosine_similarity(f, b, dim=-1) * self.hp.cosine_coef
+    # @profile
+    def forward(self, spect, seg, phonemes, length):
+        device = next(self.parameters()).device
+        spect = spect.to(device)
+        if length is not None and isinstance(length, torch.Tensor):
+            length = length.to(device)
+        z = self.enc(spect.unsqueeze(1))
+        
+        del spect
+        torch.cuda.empty_cache()
+        
+        z = F.normalize(z, dim=-1)
+
+        z_bilstm, _ = self.bi_lstm(z)
+        logits = self.fc(z_bilstm)
+        probs = F.softmax(logits, dim=-1)
+        
+        probs = logits        
+        frame_labels = torch.zeros_like(logits, device=device)
+        
+        preds = defaultdict(list)
+        for i, t in enumerate(self.pred_steps):
+            positive_b_scores_list = []
+            negative_b_scores_list = []
+            if seg is None:
+                z_seg = z
+                score = self.score(z_seg[:, :-t], z_seg[:, t:])
+                for b in range(z.shape[0]):
+                    positive_b_scores_list.append(score[b])
+                    negative_b_scores_list.append(score[b])
+            else:
+                for b in range(z.shape[0]):
+                    seg_b = [0] + seg[b][:] + [z[b].shape[0]]
+                    pos_scores = []
+                    neg_scores = []
+                    segment_pairs = list(zip(seg_b[:-1], seg_b[1:-1]))
+                    num_repeats = 5
+                    for ki, kii in segment_pairs:
+                        mid = ki + (kii - ki) // 2
+                        sample_len = 1
+                        start_idx = int(ki + 0.2 * (kii - ki))
+                        end_idx = int(ki + 0.8 * (kii - ki))
+                        num_idx = end_idx - start_idx
+                        # Batch positive sampling
+                        if num_idx >= sample_len:
+                            idxs = torch.tensor(
+                                random.choices(range(start_idx, end_idx), k=num_repeats),
+                                device=device
+                            )
+                            z_pos = z[b, idxs]
+                        else:
+                            idx_start = int(mid - (sample_len // 2))
+                            idx_end = int(mid + (sample_len // 2))
+                            z_pos = z[b, idx_start:idx_end].repeat(num_repeats, 1)
+                        pos_scores.append(self.score(z_pos[:, :-t], z_pos[:, t:]))
+                        # Batch negative sampling
+                        if ki - (sample_len // 2) <= 0 or kii + (sample_len // 2) >= z[b].shape[0]:
+                            idx_start = int(mid - (sample_len // 2))
+                            idx_end = int(mid + (sample_len // 2))
+                            z_neg = z[b, idx_start:idx_end].repeat(num_repeats, 1)
+                        else:
+                            idx_start = int(kii - (sample_len // 2))
+                            idx_end = int(kii + (sample_len // 2))
+                            z_neg = z[b, idx_start:idx_end].repeat(num_repeats, 1)
+                        neg_scores.append(self.score(z_neg[:, :-t], z_neg[:, t:]))
+                    if pos_scores:
+                        positive_b_scores_list.append(torch.cat(pos_scores))
+                        negative_b_scores_list.append(torch.cat(neg_scores))
+            # Padding and stacking
+            if positive_b_scores_list:
+                original_lengths = torch.tensor([x.shape[0] for x in positive_b_scores_list], dtype=torch.int64, device=device)
+                max_length = original_lengths.max().item()
+                def pad_scores(scores):
+                    padded = torch.full((max_length,), float(0), device=device)
+                    padded[:scores.shape[0]] = scores
+                    return padded
+                positive_b_scores = torch.stack([pad_scores(x) for x in positive_b_scores_list])
+                negative_b_scores = torch.stack([pad_scores(x) for x in negative_b_scores_list])
+            else:
+                batch_size = z.shape[0]
+                max_length = 1
+                positive_b_scores = torch.full((batch_size, max_length), float(0), device=device)
+                negative_b_scores = torch.full((batch_size, max_length), float(0), device=device)
+                original_lengths = torch.ones(batch_size, dtype=torch.int64, device=device)
+            preds[t].append(positive_b_scores)
+            preds[t].append(negative_b_scores)
+
+        # Peak detection and post-processing (not differentiable, but kept for completeness)
+        total_peaks = []
+        for i in range(len(phonemes)):
+            if seg is not None:
+                segments = seg[i]
+            else:
+                segments = []
+
+            probs_real = F.softmax(probs[i], dim=-1).squeeze(0)        
+            
+            cur_preds = preds[1][0][i]
+            cur_preds = max_min_norm(cur_preds)
+            
+            median_h = cur_preds.median()
+            preds_np = cur_preds - median_h
+            w_phi = torch.softmax(self.w_phi, dim=0)
+            
+
+            sr = 16000
+            len_ratio = 161.34011627906978
+            preds_peaks = detect_peaks(
+                x= (cur_preds), #(-1 * cur_preds),
+                w_phi=w_phi,
+                original_lengths_all=original_lengths[i],
+                phonemes=phonemes[i],
+                len_ratio=len_ratio,
+                probs_real_all=probs_real
+            )
+            preds_peaks = preds_peaks[0] * len_ratio / sr
+            if seg is not None:
+                segments = np.array(segments) * len_ratio / sr
+
+            # print("truth boundaries ('segments'):")
+            # print(segments)
+            # print("predicted boundaries (in seconds):")
+            # print(preds_peaks)
+
+            total_peaks.append(preds_peaks)
+
+        # for debug mode plotting only ---> probs=z
+        return preds, original_lengths, probs, frame_labels, seg, total_peaks, w_phi
+
+
+    def loss_ph(self, preds, original_lengths, probs, frame_labels, seg, total_peaks, w_phi, phonemes):
+        if seg is not None:
+            for i, (segments, phs) in enumerate(zip(seg, phonemes)):
+                starts = np.array([0] + list(segments[:-1]), dtype=int)
+                ends = np.array(segments[:], dtype=int)
+                labels = [timit_to_leehon(p) or 'sil' for p in phs] 
+                label_indices = [self.phoneme_to_index[l] for l in labels]
+                for start, end, label_index in zip(starts, ends, label_indices):
+                    frame_labels[i, start:end, label_index] = 1.0 
+        total_loss = 0.0
+        probs = probs.view(-1, probs.size(-1))
+        frame_labels = frame_labels.view(-1, frame_labels.size(-1))
+        ph_loss = F.cross_entropy(probs, frame_labels.argmax(dim=-1))
+        loss = 0
+        
+        for t, t_preds in preds.items():
+            mask = length_to_mask(original_lengths - t + 1)
+            out = torch.stack(t_preds, dim=-1)
+            out = F.log_softmax(out, dim=-1)
+            pos_loss = out[..., 0] * mask
+            neg_loss = out[..., 1] * mask
+
+            pos_count = mask.sum()
+            neg_count = mask.sum()
+            
+            pos_loss_mean = pos_loss.sum() / pos_count if pos_count > 0 else torch.tensor(0.0, device=pos_loss.device)
+            neg_loss_mean = neg_loss.sum() / neg_count if neg_count > 0 else torch.tensor(0.0, device=neg_loss.device)
+
+            w_pos_neg = torch.softmax(self.w_pos_neg, dim=0)
+
+            l_pos_neg = torch.stack([pos_loss_mean, -neg_loss_mean])
+            loss += -torch.dot(w_pos_neg, l_pos_neg)
+
+        total_loss = (loss) + 0.2*ph_loss
+
+        sum_mse = torch.tensor(0.0, device=probs.device)
+
+        if seg is not None and len(total_peaks) > 0:
+            seg_tensors = [torch.tensor(s, dtype=torch.float32, device=probs.device) for s in seg]
+            for i in range(len(seg_tensors)):
+                seg_tensors[i] = seg_tensors[i]*161.34011627906978/16000  # convert to seconds
+            peaks_tensors = [torch.tensor(tp[1:], dtype=torch.float32, device=probs.device) for tp in total_peaks]
+
+            mse = 0.0
+            for i in range(len(seg_tensors)):
+                mse = mse + F.mse_loss(seg_tensors[i], peaks_tensors[i])
+            mse = mse / len(seg_tensors)
+
+            sum_mse = mse
+
+            total_loss = ph_loss
+
+        return total_loss, ph_loss, loss, sum_mse, w_pos_neg, w_phi
+    
+    def loss_nce(self, preds, original_lengths, probs, frame_labels, seg, total_peaks, w_phi, phonemes):
+        if seg is not None:
+            for i, (segments, phs) in enumerate(zip(seg, phonemes)):
+                starts = np.array([0] + list(segments[:-1]), dtype=int)
+                ends = np.array(segments[:], dtype=int)
+                labels = [timit_to_leehon(p) or 'sil' for p in phs]
+                label_indices = [self.phoneme_to_index[l] for l in labels]
+                for start, end, label_index in zip(starts, ends, label_indices):
+                    frame_labels[i, start:end, label_index] = 1.0
+        
+        total_loss = 0.0
+        probs = probs.view(-1, probs.size(-1))
+        frame_labels = frame_labels.view(-1, frame_labels.size(-1))
+        ph_loss = F.cross_entropy(probs, frame_labels.argmax(dim=-1))
+        loss = 0
+        
+        for t, t_preds in preds.items():
+            mask = length_to_mask(original_lengths - t + 1)
+            out = torch.stack(t_preds, dim=-1)
+            out = F.log_softmax(out, dim=-1)
+            pos_loss = out[..., 0] * mask
+            neg_loss = out[..., 1] * mask
+
+            pos_count = mask.sum()
+            neg_count = mask.sum()
+            
+            pos_loss_mean = pos_loss.sum() / pos_count if pos_count > 0 else torch.tensor(0.0, device=pos_loss.device)
+            neg_loss_mean = neg_loss.sum() / neg_count if neg_count > 0 else torch.tensor(0.0, device=neg_loss.device)
+
+            w_pos_neg = torch.softmax(self.w_pos_neg, dim=0)
+
+            l_pos_neg = torch.stack([pos_loss_mean, -neg_loss_mean])
+            loss += -torch.dot(w_pos_neg, l_pos_neg)
+
+        total_loss = (loss)
+
+        sum_mse = torch.tensor(0.0, device=probs.device)
+
+        if seg is not None and len(total_peaks) > 0:
+            seg_tensors = [torch.tensor(s, dtype=torch.float32, device=probs.device) for s in seg]
+            for i in range(len(seg_tensors)):
+                seg_tensors[i] = seg_tensors[i]*161.34011627906978/16000  # convert to seconds
+            peaks_tensors = [torch.tensor(tp[1:], dtype=torch.float32, device=probs.device) for tp in total_peaks]
+
+            mse = 0.0
+            for i in range(len(seg_tensors)):
+                mse = mse + F.mse_loss(seg_tensors[i], peaks_tensors[i])
+            mse = mse / len(seg_tensors)
+
+            sum_mse = mse
+            total_loss = loss
+
+        return total_loss, ph_loss, loss, sum_mse, w_pos_neg, w_phi
+    
+    # ------------------------------ For Ablations ---------------------------------
+    def loss_InfoNCE_classic(self, preds, original_lengths, probs, frame_labels, seg, total_peaks, w_phi, phonemes):
+        """
+        Classic InfoNCE implementation.
+        Standard log-softmax over one positive and N negatives.
+        """
+        device = probs.device
+        total_loss = 0.0
+        probs_flat = probs.view(-1, probs.size(-1))
+        frame_labels_flat = frame_labels.view(-1, frame_labels.size(-1))
+        ph_loss = F.cross_entropy(probs_flat, frame_labels_flat.argmax(dim=-1))
+
+        nce_loss_accum = 0.0
+        
+        for t, t_preds in preds.items():
+            logits = torch.stack(t_preds, dim=-1) 
+            batch_size, seq_len, _ = logits.shape
+            target = torch.zeros((batch_size, seq_len), dtype=torch.long, device=device)
+            mask = length_to_mask(original_lengths - t + 1)
+            logits_flat = logits.view(-1, 2)
+            target_flat = target.view(-1)
+            loss_fn = nn.CrossEntropyLoss(reduction='none')
+            raw_loss = loss_fn(logits_flat, target_flat)
+            masked_loss = raw_loss * mask.view(-1)
+            nce_loss_accum += masked_loss.sum() / mask.sum()
+        total_loss = nce_loss_accum / len(preds)
+        sum_mse = torch.tensor(0.0, device=device)
+        w_pos_neg = torch.softmax(self.w_pos_neg, dim=0) 
+
+        return total_loss, ph_loss, nce_loss_accum, sum_mse, w_pos_neg, w_phi
+    # -----------------------------------------------------------------------------
+    
+    def loss_mse(self, preds, original_lengths, probs, frame_labels, seg, total_peaks, w_phi, phonemes):
+        if seg is not None:
+            for i, (segments, phs) in enumerate(zip(seg, phonemes)):
+                starts = np.array([0] + list(segments[:-1]), dtype=int)
+                ends = np.array(segments[:], dtype=int)
+                labels = [timit_to_leehon(p) or 'sil' for p in phs]
+                label_indices = [self.phoneme_to_index[l] for l in labels]
+                for start, end, label_index in zip(starts, ends, label_indices):
+                    frame_labels[i, start:end, label_index] = 1.0
+        
+        total_loss = 0.0
+        probs = probs.view(-1, probs.size(-1))
+        frame_labels = frame_labels.view(-1, frame_labels.size(-1))
+        ph_loss = F.cross_entropy(probs, frame_labels.argmax(dim=-1))
+        loss = 0
+        
+        for t, t_preds in preds.items():
+            mask = length_to_mask(original_lengths - t + 1)
+            out = torch.stack(t_preds, dim=-1)
+            out = F.log_softmax(out, dim=-1)
+            pos_loss = out[..., 0] * mask
+            neg_loss = out[..., 1] * mask
+
+            pos_count = mask.sum()
+            neg_count = mask.sum()
+            
+            pos_loss_mean = pos_loss.sum() / pos_count if pos_count > 0 else torch.tensor(0.0, device=pos_loss.device)
+            neg_loss_mean = neg_loss.sum() / neg_count if neg_count > 0 else torch.tensor(0.0, device=neg_loss.device)
+
+            w_pos_neg = torch.softmax(self.w_pos_neg, dim=0)
+
+            l_pos_neg = torch.stack([pos_loss_mean, -neg_loss_mean])
+            loss += -torch.dot(w_pos_neg, l_pos_neg)
+
+        total_loss = (loss) + 0.2*ph_loss
+        sum_mse = torch.tensor(0.0, device=probs.device)
+
+        if seg is not None and len(total_peaks) > 0:
+            seg_tensors = [torch.tensor(s, dtype=torch.float32, device=probs.device) for s in seg]
+            for i in range(len(seg_tensors)):
+                seg_tensors[i] = seg_tensors[i]*161.34011627906978/16000  # convert to seconds
+            peaks_tensors = [torch.tensor(tp[1:], dtype=torch.float32, device=probs.device) for tp in total_peaks]
+            mse = 0.0
+            for i in range(len(seg_tensors)):
+                mse = mse + F.mse_loss(seg_tensors[i], peaks_tensors[i])
+            mse = mse / len(seg_tensors)
+            sum_mse = mse
+            total_loss = sum_mse
+
+        return total_loss, ph_loss, loss, sum_mse, w_pos_neg, w_phi
+
+    def total_loss(self, preds, original_lengths, probs, frame_labels, seg, total_peaks, w_phi,phonemes):
+        if seg is not None:
+            for i, (segments, phs) in enumerate(zip(seg, phonemes)):
+                starts = np.array([0] + list(segments[:-1]), dtype=int)
+                ends = np.array(segments[:], dtype=int)
+                labels = [timit_to_leehon(p) or 'sil' for p in phs]
+                label_indices = [self.phoneme_to_index[l] for l in labels]
+                for start, end, label_index in zip(starts, ends, label_indices):
+                    frame_labels[i, start:end, label_index] = 1.0
+        
+        total_loss = 0.0
+        probs = probs.view(-1, probs.size(-1))
+        frame_labels = frame_labels.view(-1, frame_labels.size(-1))
+        ph_loss = F.cross_entropy(probs, frame_labels.argmax(dim=-1))
+        loss = 0
+        
+        for t, t_preds in preds.items():
+            mask = length_to_mask(original_lengths - t + 1)
+            out = torch.stack(t_preds, dim=-1)
+            out = F.log_softmax(out, dim=-1)
+            pos_loss = out[..., 0] * mask
+            neg_loss = out[..., 1] * mask
+
+            pos_count = mask.sum()
+            neg_count = mask.sum()
+            
+            pos_loss_mean = pos_loss.sum() / pos_count if pos_count > 0 else torch.tensor(0.0, device=pos_loss.device)
+            neg_loss_mean = neg_loss.sum() / neg_count if neg_count > 0 else torch.tensor(0.0, device=neg_loss.device)
+
+            w_pos_neg = torch.softmax(self.w_pos_neg, dim=0)
+
+            l_pos_neg = torch.stack([pos_loss_mean, -neg_loss_mean])
+            loss += -torch.dot(w_pos_neg, l_pos_neg)
+
+        total_loss = (loss) + 0.2*ph_loss
+        sum_mse = torch.tensor(0.0, device=probs.device)
+
+        if seg is not None and len(total_peaks) > 0:
+            seg_tensors = [torch.tensor(s, dtype=torch.float32, device=probs.device) for s in seg]
+            for i in range(len(seg_tensors)):
+                seg_tensors[i] = seg_tensors[i]*161.34011627906978/16000  # convert to seconds
+            peaks_tensors = [torch.tensor(tp[1:], dtype=torch.float32, device=probs.device) for tp in total_peaks]
+
+            mse = 0.0
+            for i in range(len(seg_tensors)):
+                mse = mse + F.mse_loss(seg_tensors[i], peaks_tensors[i])
+            mse = mse / len(seg_tensors)
+
+            sum_mse = mse
+            total_loss = total_loss + sum_mse 
+        return total_loss, ph_loss, loss, sum_mse, w_pos_neg, w_phi
+
+@hydra.main(config_path='conf/config.yaml', strict=False)
+def main(cfg):
+    ds, _, _ = TrainTestDataset.get_datasets(cfg.timit_path)
+    spect, seg, phonemes, length, fname = ds[0]
+    spect = spect.unsqueeze(0)
+    model = NextFrameClassifier(cfg)
+    out = model(spect, seg, phonemes, length)
+
+if __name__ == "__main__":
+    main()

packages.txt

@@ -0,0 +1,2 @@
+espeak
+espeak-ng

predict.py

@@ -0,0 +1,293 @@
+import argparse
+from glob import glob
+from unittest import case
+import dill
+from argparse import Namespace
+import torch
+import torchaudio
+import torch.nn.functional as F
+from utils import (max_min_norm,
+                   get_timit_61_phoneme_mappings)
+from next_frame_classifier import NextFrameClassifier
+
+from dataloader import spectral_size
+import matplotlib.pyplot as plt
+import numpy as np
+import os
+
+import dutch_preprocess
+from utils import timit_to_leehon_map_MACRO, timit_leehon_39_phonemes, timit_61_phonemes
+
+# Cache the loaded model+peak-params by checkpoint path so batch runs (many files,
+# one checkpoint) don't reload ~330MB from disk per file. Output is unchanged.
+_MODEL_CACHE = {}
+
+def _load_model(ckpt_path):
+    cached = _MODEL_CACHE.get(ckpt_path)
+    if cached is not None:
+        return cached
+    ckpt = torch.load(ckpt_path, map_location=lambda storage, loc: storage)
+    hp = ckpt["hparams"]
+    model = NextFrameClassifier(hp)
+    try:
+        weights = ckpt["state_dict"]
+    except Exception:
+        weights = ckpt["model_state_dict"]
+    weights = {k.replace("NFC.", ""): v for k, v in weights.items()}
+    model.load_state_dict(weights)
+    model.eval()
+    peak_detection_params = dill.loads(ckpt['peak_detection_params'])['cpc_1']
+    _MODEL_CACHE.clear()  # keep only the most-recent checkpoint (bounded memory)
+    _MODEL_CACHE[ckpt_path] = (model, peak_detection_params)
+    return model, peak_detection_params
+
+def main_predict(wav, ckpt, w_phi, language="english", annotation="phn", no_plots=False):
+    print(f"running inference on: {wav}")
+    print(f"running inferece using ckpt: {ckpt}")
+    print("\n\n", 90 * "-")
+
+    # Optional plot suppression for fast batch runs. Patches plt.savefig for
+    # the duration of this call so plots inside utils.phoneme_alignment are
+    # skipped too, without having to modify utils.py.
+    _orig_savefig = plt.savefig if no_plots else None
+    if no_plots:
+        plt.savefig = lambda *a, **k: None
+
+    model, peak_detection_params = _load_model(ckpt)
+    # peak_detection_params["prominence"] = prominence # Unused
+    # load data
+    audio, sr = torchaudio.load(wav)
+    assert sr == 16000, "model was trained with audio sampled at 16khz, please downsample."
+    audio = audio[0]
+    # audio = audio.unsqueeze(0)
+    
+    base_dir = os.path.dirname(wav)
+    base_name = os.path.basename(wav).split('.')[0]
+    # search_pattern = os.path.join(base_dir, f"{base_name}*.phn")
+    # search_pattern = os.path.join(base_dir, f"{base_name}*.wrd")
+    # search_pattern = os.path.join(base_dir, f"{base_name}*.word")
+    search_pattern = os.path.join(base_dir, f"{base_name}*.{annotation}")
+    matching_files = glob(search_pattern)
+    if matching_files:
+        phn_path = matching_files[0]
+    else:
+        print("No matching .phn file found. Using default naming convention.")
+        phn_path = wav.replace("wav", "phn")
+
+    # load audio
+    audio_len = len(audio)
+    spectral_len = spectral_size(audio_len)
+    len_ratio = (audio_len / spectral_len)
+
+    # load labels -- segmentation and phonemes
+    with open(phn_path, "r") as f:
+        lines = f.readlines()
+        lines = list(map(lambda line: line.split(), lines))
+
+        # get segment times
+        times = torch.FloatTensor(list(map(lambda line: int(float(line[1]) / len_ratio), lines)))[:-1]  # don't count end time as boundary
+        # times = torch.FloatTensor(list(map(lambda line: int(int(line[1]) / len_ratio), lines)))[:-1]  # don't count end time as boundary
+        times_sec = torch.FloatTensor(list(map(lambda line: (float(line[1]) / sr), lines)))[:-1]  # don't count end #sr = 16000 in TIMIT
+        # times_sec = torch.FloatTensor(list(map(lambda line: (int(line[1]) / sr), lines)))[:-1]  # don't count end #sr = 16000 in TIMIT
+        
+        # get phonemes in each segment (for K times there should be K+1 phonemes)
+        phonemes = list(map(lambda line: line[2].strip(), lines))
+
+    # Original input labels (one per .phn/.wrd/.txt line) — kept around for the
+    # truth-boundary text on plots. After G2P the `phonemes` variable is
+    # flattened into LH39 phonemes whose count differs from len(times), so we
+    # cache the per-line labels here before that transformation.
+    truth_labels = list(phonemes)
+    
+    if language == "dutch":
+        lh39_ph = []
+        for IFA_ph in phonemes:
+            print(f"\nINPUT: {IFA_ph}")
+            # output = dutch_preprocess.aligner_pipeline(timit_to_leehon_map_MACRO[IFA_ph.lower()])
+            output = dutch_preprocess.aligner_pipeline(IFA_ph if IFA_ph.lower() not in timit_61_phonemes else timit_to_leehon_map_MACRO[IFA_ph.lower()])
+            # output = dutch_preprocess.aligner_pipeline(IFA_ph)
+            
+            # # FOR WORDS: #not good
+            # lh39_ph.append(output[0]["lh39"])
+            
+            # FOR PHONEMES:
+            lh39_ph.append([x["lh39"] for x in output])
+        if not output:
+            print("Results: None")
+        print(phonemes)
+        print(f"Dutch IPA to LH39 mapping: {lh39_ph}")
+        
+        # try
+        phonemes = np.hstack(lh39_ph).tolist()
+        
+    audio, seg, phonemes, length = audio.unsqueeze(0), [times.tolist()], [phonemes], [audio_len/len_ratio] #[spectral_size(len(audio))]
+
+    
+    
+    with torch.no_grad():
+        model.eval()
+        
+        # preds,original_lengths, probs, frame_labels = model(audio,None,phonemes,length)
+        # ------- Sept 10 - check with truth preds no truth for nce ---------------
+        preds,original_lengths, probs, frame_labels, _,preds_peaks, w_phi = model(audio,None,phonemes,length)
+        # preds,original_lengths, probs, frame_labels, _,preds_peaks, w_phi = model(audio,seg,phonemes,length)
+        # ------- ------------------------------------------------- ---------------
+    
+    phoneme_to_idx, idx_to_phoneme = get_timit_61_phoneme_mappings()
+    
+    
+    phoneme_labels = [idx_to_phoneme[i] for i in range(39)]
+    # phoneme_labels = [idx_to_phoneme[i] for i in range(61)]
+    # phoneme_labels = [idx_to_phoneme[i] for i in range(41)]
+    # probs_real = probs #F.softmax(probs, dim=-1) 
+    probs_real = F.softmax(probs, dim=-1).squeeze(0).detach().numpy()
+    
+    out_dir = os.path.dirname(wav)
+    base_name = os.path.basename(wav).replace('.wav', '')
+
+    plt.figure(figsize=(15, 5))
+    plt.imshow(probs_real.T, aspect='auto', cmap='viridis')
+    plt.colorbar(label='Probability')
+    plt.xlabel('Frame Index')
+    plt.ylabel('Phoneme')
+    plt.yticks(ticks=range(39), labels=phoneme_labels)
+    plt.title('Frame-wise Label Probability Map')
+    # Truth-boundary axvlines + per-segment label text. truth_labels is the
+    # original (pre-G2P) per-line input labels, so its length matches `times`.
+    for i, s in enumerate(times):
+        s_val = float(s)
+        plt.axvline(x=s_val, color='red', linestyle='--', linewidth=1,
+                    label='Truth boundary' if i == 0 else "")
+        if i < len(truth_labels):
+            plt.text(s_val, probs_real.shape[1] + 1, truth_labels[i],
+                     color='red', rotation=90, va='top', ha='center', fontsize=8)
+    plt.savefig(os.path.join(out_dir, f"{base_name}_probs.png"))
+    plt.close()
+    
+    # Latent boundary scores (CNN output) — used for the boundaries plot.
+    preds = preds[1][0]
+    preds = max_min_norm(preds)
+    preds_np = preds.detach().numpy()[0]
+    median_h = np.median(preds_np)
+    preds_np = preds_np - median_h
+
+    # Predicted boundary timestamps (in seconds).
+    preds = torch.tensor(preds_peaks[0], dtype=torch.float32)
+    print(f"predicted boundaries (s): {preds}")
+
+    # Marker height: scale to the *typical* peak of the latent score so the
+    # triangles sit roughly on top of peaks instead of dwarfing them or
+    # vanishing into the noise floor. The signal is sparse, so we filter out
+    # near-zero values and take the median of what's left — robust to both
+    # outlier spikes and the long tail of zero-ish frames.
+    abs_p = np.abs(preds_np) if preds_np.size else np.array([])
+    if abs_p.size:
+        noise_thresh = 0.05 * float(np.max(abs_p))
+        peaks = abs_p[abs_p > noise_thresh]
+        marker_h = float(np.median(peaks)) if peaks.size else float(np.max(abs_p))
+    else:
+        marker_h = 0.1
+    if marker_h <= 0:
+        marker_h = 0.1
+
+    # Truth-boundary marker signal (peak per truth boundary, zero elsewhere).
+    signal = np.zeros(int(original_lengths[0]))
+    signal_max_idx = signal.shape[0] - 1
+    for t in times:
+        idx = min(int(t), signal_max_idx)
+        signal[idx] = marker_h
+    times_clipped = [t for t in times if float(t) <= signal_max_idx]
+
+    plt.figure(figsize=(12, 6))
+    plt.plot(signal, marker='*', linestyle='-', label='Truth boundary marker')
+    derivative_preds_np = np.diff(preds_np)
+    derivative_preds_np = np.concatenate([[0], derivative_preds_np])
+    plt.plot(range(len(derivative_preds_np)), derivative_preds_np,
+             marker='o', label='Derivative of latent score', color='magenta')
+    plt.plot(range(len(preds_np)), preds_np,
+             marker='*', label='Latent score', color='red')
+
+    preds_plot = np.zeros(int(original_lengths[0]))
+    for pred in preds:
+        idx = int(pred * sr / len_ratio)
+        if idx <= len(preds_plot) - 1:
+            preds_plot[idx] = marker_h
+    plt.plot(range(len(preds_plot)), preds_plot,
+             marker='^', label='Predicted boundary', linestyle='None')
+
+    y_top = plt.ylim()[1]
+    for i, s in enumerate(times_clipped):
+        s_val = float(s)
+        plt.axvline(x=s_val, color='red', linestyle='--', linewidth=1,
+                    label='Truth boundary' if i == 0 else "")
+        if i < len(truth_labels):
+            plt.text(s_val, y_top, truth_labels[i], color='red',
+                     rotation=90, va='top', ha='center', fontsize=8)
+
+    plt.xlabel('Frame Index')
+    plt.ylabel('Score')
+    plt.title('Predicted Boundaries')
+    plt.legend(loc='upper right')
+    # Clip the y-axis to a few × typical-peak height so single outlier spikes
+    # don't visually squash the rest of the signal.
+    y_half = max(marker_h * 4.0, 0.1)
+    plt.ylim(-y_half, y_half)
+    plt.savefig(os.path.join(out_dir, f"{base_name}_boundaries.png"))
+    plt.close()
+
+    pred_bound, truth_bound = preds, times_sec
+    mapped_ph = lh39_ph if language == "dutch" else None
+
+    if no_plots:
+        plt.savefig = _orig_savefig
+
+    return pred_bound, truth_bound, mapped_ph
+
+_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+DEFAULT_CKPT_ENGLISH      = os.path.join(_SCRIPT_DIR, "pretrained_models", "falcon_timit_english.pt")
+# The joint TIMIT+Buckeye model generalizes best to unseen languages — it is the
+# strongest checkpoint on every multilingual test set (Dutch/German/Hebrew,
+# phoneme and word) — so it is the default for the multilingual / cross-lingual path.
+DEFAULT_CKPT_MULTILINGUAL = os.path.join(_SCRIPT_DIR, "pretrained_models", "falcon_joint_multilingual.pt")
+
+def resolve_internal_language(lang: str, mode: str, annotation: str) -> str:
+    """
+    Map user-facing (--lang, --mode, --annotation) to the internal `language`
+    flag main_predict() understands.
+
+      'english' = no G2P; assumes labels are already TIMIT-39 phonemes.
+      'dutch'   = G2P pipeline (panphon-based mapping). Used for any non-English
+                  language, word-level mode, or plain-text input.
+    """
+    if lang == "english" and mode == "phoneme" and annotation.lower() == "phn":
+        return "english"
+    return "dutch"
+
+def resolve_default_ckpt(lang: str) -> str:
+    return DEFAULT_CKPT_MULTILINGUAL if lang == "multilingual" else DEFAULT_CKPT_ENGLISH
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Unsupervised segmentation inference script')
+    parser.add_argument('--wav', help='path to wav file')
+    parser.add_argument('--ckpt', default=None,
+                        help='Path to checkpoint file. If omitted, uses '
+                             'pretrained_models/falcon_timit_english.pt for --lang english '
+                             'or falcon_joint_multilingual.pt for --lang multilingual '
+                             '(the joint TIMIT+Buckeye model is best for cross-lingual zero-shot).')
+
+    parser.add_argument('--mode', type=str, default='phoneme', choices=['phoneme', 'word'],
+                        help='Alignment granularity: "phoneme" = phoneme-level alignment (default). '
+                             '"word" = word-level alignment (zero-shot, no additional training).')
+    parser.add_argument('--lang', type=str, default='english', choices=['english', 'multilingual'],
+                        help='Language setting: "english" = trained English phoneme alignment (default). '
+                             '"multilingual" = any non-English language (zero-shot cross-lingual).')
+    parser.add_argument('--annotation', type=str, default='phn',
+                        help='Annotation file extension to search for (e.g. phn, wrd, word, txt). Default: phn')
+    parser.add_argument('--no-plots', action='store_true',
+                        help='Skip per-file diagnostic plots (probs/logits/boundaries/dp_matrix) for faster runs.')
+    args = parser.parse_args()
+
+    ckpt = args.ckpt or resolve_default_ckpt(args.lang)
+    language = resolve_internal_language(args.lang, args.mode, args.annotation)
+    main_predict(args.wav, ckpt, w_phi=0.5, language=language,
+                 annotation=args.annotation, no_plots=args.no_plots)

requirements.txt

@@ -0,0 +1,22 @@
+# FALCON Space — pinned to the working local env (Python 3.8) + pynini for MFA-like G2P
+torch==2.4.1
+torchaudio==2.4.1
+hydra-core==0.11.3
+omegaconf==1.4.1
+wandb==0.8.24
+torch-optimizer==0.0.1a12
+librosa==0.8.1
+soundfile==0.13.1
+numpy==1.22.4
+scipy==1.7.3
+scikit-learn==1.3.2
+panphon==0.21.0
+dill==0.3.1.1
+boltons==20.0.0
+matplotlib==3.4.3
+memory-profiler==0.61.0
+tqdm==4.62.3
+gradio==4.44.1
+huggingface_hub==0.36.2
+TextGrid>=1.5
+pynini==2.1.6

utils.py

@@ -0,0 +1,691 @@
+import os
+import random
+import torch
+import torch.nn as nn
+import numpy as np
+import matplotlib.pyplot as plt
+import time
+from scipy.signal import find_peaks
+import wandb
+from tqdm import tqdm
+import concurrent.futures
+from typing import List, Sequence, Union
+import time
+from memory_profiler import profile
+
+# Optionally redirect the dp_matrix plot to a specific directory (used by demo).
+# Set via set_dp_matrix_out_dir() before calling inference; reset to None after.
+_dp_matrix_out_dir = None
+
+def set_dp_matrix_out_dir(path):
+    global _dp_matrix_out_dir
+    _dp_matrix_out_dir = path
+
+timit_leehon_39_phonemes = [
+    'ao', 'ae', 'ah','aw', 'er', 'ay', 
+    'b', 'sil', 'ch', 'd', 'dh', 'dx', 'eh', 'el', 'm', 'en', 'ng', 'ey',
+    'f', 'g', 'hh', 'ih', 'iy', 'jh', 'k', 'v', 'w', 'y', 'z', 'sh', 't', 'r', 's', 'th','uh', 'uw', 'oy', 'ow','p'
+]
+
+timit_61_phonemes = [
+        'aa', 'ae', 'ah', 'ao', 'aw', 'ax', 'ax-h', 'axr', 'ay',
+        'b', 'bcl', 'ch', 'd', 'dcl', 'dh', 'dx', 'eh', 'el', 'em', 'en', 'eng', 'epi', 'er', 'ey',
+        'f', 'g', 'gcl', 'h#', 'hh', 'hv', 'ih', 'ix', 'iy', 'jh', 'k', 'kcl', 'm', 'n', 'ng', 'l',
+        'nx', 'ow', 'oy', 'p', 'pau', 'pcl', 'q', 'r', 's', 'sh', 't', 'tcl', 'th', 'uh', 'uw','ux',
+        'v', 'w', 'y', 'z', 'zh'
+    ]
+
+# Create mappings
+# phoneme_to_idx = {phoneme: idx for idx, phoneme in enumerate(timit_61_phonemes)}
+phoneme_to_idx_MACRO = {phoneme: idx for idx, phoneme in enumerate(timit_leehon_39_phonemes)}
+
+idx_to_phoneme_MACRO = {idx: phoneme for phoneme, idx in phoneme_to_idx_MACRO.items()}
+timit_to_leehon_map_MACRO = {
+        'aa': 'ao', 'ae': 'ae', 'ah': 'ah', 'ao': 'ao', 'aw': 'aw', 'ax': 'ah', 'ax-h': 'ah', 'axr': 'er', 'ay': 'ay',
+        'b': 'b', 'bcl': 'sil', 'ch': 'ch', 'd': 'd', 'dcl': 'sil', 'dh': 'dh', 'dx': 'dx', 'eh': 'eh', 'el': 'el',
+        'em': 'm', 'en': 'en', 'eng': 'ng', 'epi': 'sil', 'er': 'er', 'ey': 'ey', 'f': 'f', 'g': 'g', 'gcl': 'sil',
+        'h#': 'sil', 'hh': 'hh', 'hv': 'hh', 'ih': 'ih', 'ix': 'ih', 'iy': 'iy', 'jh': 'jh', 'k': 'k', 'kcl': 'sil',
+        'l': 'el', 'm': 'm', 'n': 'en', 'ng': 'ng', 'nx': 'en', 'ow': 'ow', 'oy': 'oy', 'p': 'p', 'pau': 'sil', 'pcl': 'sil',
+        'q': 't', 'qcl': 'sil', 'r': 'r', 's': 's', 'sh': 'sh', 't': 't', 'tcl': 'sil', 'th': 'th', 'uh': 'uh', 'uw': 'uw',
+        'ux': 'uw', 'v': 'v', 'w': 'w', 'y': 'y', 'z': 'z', 'zh': 'sh',
+}
+
+def create_truth_probs_real(segments, phonemes, phoneme_to_index, num_frames):
+    segments = [0] + list(segments)
+    num_phonemes = len(phoneme_to_index)
+    probs_real = torch.zeros((num_frames, num_phonemes), dtype=torch.float32)
+    for seg_idx in range(len(phonemes)):
+        start = int(segments[seg_idx])
+        end = int(segments[seg_idx + 1]) if seg_idx + 1 < len(segments) else num_frames
+        if end > start:
+            ph_label = phonemes[seg_idx].lower()
+            ph_index = phoneme_to_index.get(ph_label, phoneme_to_index.get('sil', 0))
+            probs_real[start:end, ph_index] = 1.0
+    return probs_real
+
+
+# ------------------------------- ablations -------------------------------
+
+# phoneme alignment with classic (hard) DP
+def phoneme_alignment_Hard_DP(p_seq, w_phi, original_lengths, len_ratio, derivative_preds_np, probs_real):
+    # Gamma is kept for signature consistency but not used in hard DP
+    gamma = 1e-20 
+    T = int(original_lengths[0])
+    n = len(p_seq)
+    device = derivative_preds_np.device
+
+    if isinstance(probs_real, np.ndarray):
+        probs_real = torch.tensor(probs_real, device=device)
+    cumsum_probs = torch.cumsum(probs_real, dim=0)
+    
+    phoneme_mappings = {p.lower(): timit_to_leehon_map_MACRO.get(p.lower(), 'sil') if p.lower() not in timit_leehon_39_phonemes else p.lower() for p in p_seq}
+    derivatives = torch.cat([torch.tensor([0], device=device), torch.diff(derivative_preds_np, dim=0)])
+
+    # Initialize DP matrix with very low value
+    dp_mat = torch.full((n, T, T), float(-1e9), device=device)
+    p_idx0 = phoneme_to_idx_MACRO[phoneme_mappings[p_seq[0].lower()]]
+
+    # Initial state for first phoneme
+    t_e = torch.arange(T, device=device)
+    dp_mat[0, 0, :] = (
+        w_phi[0] * compute_phi_1(derivatives, 0, t_e)
+    )
+
+    # Forward Pass
+    for i in tqdm(range(1, n)):
+        p_idx = phoneme_to_idx_MACRO[phoneme_mappings[p_seq[i].lower()]]
+        t_start = torch.arange(T, device=device)
+        t_end = torch.arange(T, device=device)
+        t_start_grid, t_end_grid = torch.meshgrid(t_start, t_end, indexing='ij')
+        valid_mask = t_start_grid < t_end_grid
+        
+        phi1_dev = compute_phi_1(derivatives, t_start_grid, t_end_grid)
+        phi2 = compute_phi_2(cumsum_probs, p_idx, t_start_grid, t_end_grid)
+        total_phi = w_phi[0] * phi1_dev
+        
+        prev_scores = torch.full((T, T), float(-1e9), device=device)
+        
+        for t_end_val in range(T):
+            valid_starts = t_start[t_start < t_end_val]
+            if valid_starts.numel() == 0:
+                continue
+            
+            # --- CLASSIC DP CHANGE ---
+            # Instead of LogSumExp (Soft-Max), use Hard Max
+            prev = dp_mat[i-1, :valid_starts[-1]+1, valid_starts]
+            max_prev, _ = torch.max(prev, dim=0) 
+            prev_scores[valid_starts, t_end_val] = max_prev
+
+        dp_mat[i] = torch.where(valid_mask, total_phi + prev_scores, torch.full_like(total_phi, float(-1e9)))
+
+    # Backtracking (Classic Argmax)
+    best_start_times = torch.zeros((n), dtype=derivative_preds_np.dtype, device=device)
+    best_prev_t_end = T - 1
+    
+    for i in range(n):
+        cur_ph = n - 1 - i
+        # Find the exact index that gave the maximum score
+        scores = dp_mat[cur_ph, :, best_prev_t_end]
+        
+        # --- CLASSIC DP CHANGE ---
+        # Instead of expected_idx (Soft-Argmax), use Hard Argmax
+        best_t_start = torch.argmax(scores)
+        
+        best_start_times[cur_ph] = best_t_start.to(derivative_preds_np.dtype)
+        best_prev_t_end = int(best_t_start.item())
+        
+    # Visualization Code (unchanged logic, updated labels)
+    dp_mat_cpu = dp_mat.detach().cpu()
+    best_start_times_cpu = best_start_times.detach().cpu().numpy()
+    dp_to_plot = dp_mat_cpu.max(dim=1)[0].numpy()
+    masked_dp = np.ma.masked_where(dp_to_plot <= -1e8, dp_to_plot)
+
+    plt.figure(figsize=(12, 6))
+    cmap = plt.cm.viridis
+    cmap.set_bad(color='white')
+    plt.imshow(masked_dp, aspect='auto', origin='lower', cmap=cmap)
+    plt.colorbar(label='Hard DP Score')
+    plt.xlabel('End time (frame)')
+    plt.ylabel('Phoneme index')
+    plt.title('Classic (Hard) DP Matrix with Best Path')
+    plt.plot(best_start_times_cpu, range(len(best_start_times_cpu)), 'r.-', label='Argmax path')
+    plt.legend()
+    plt.tight_layout()
+    plt.savefig('dp_matrix_hard_classic.png')
+    plt.close()
+        
+    return best_start_times
+
+# ------------------second ablations - naive peak detection ------------------
+
+from scipy.signal import find_peaks
+
+def phoneme_alignment_naive_peak_detection(p_seq, w_phi, original_lengths, len_ratio, derivative_preds_np, probs_real):
+    """
+    Ablation version: Replaces DP with Naive Scipy Peak Detection.
+    """
+    gamma = 1e-20 
+    T = int(original_lengths[0])
+    n = len(p_seq)
+    device = derivative_preds_np.device
+
+    # --- Keep identical preprocessing to ensure 'plug & play' ---
+    if isinstance(probs_real, np.ndarray):
+        probs_real = torch.tensor(probs_real, device=device)
+    
+    # We don't actually need cumsum_probs or phoneme_mappings for naive peak detection,
+    # but we keep them defined to avoid any potential scope issues if you add code back.
+    cumsum_probs = torch.cumsum(probs_real, dim=0)
+    signal = derivative_preds_np.detach().cpu().numpy().flatten()
+
+    # --- Naive Peak Detection ---
+    # To get exactly 'n' boundaries for 'n' phonemes, we pick the top n most prominent peaks.
+    peaks, properties = find_peaks(signal, prominence=0.05) 
+    
+    peak_heights = signal[peaks]
+    
+    # Sort peaks by height and take the top 'n'
+    top_indices = np.argsort(peak_heights)[-n:]
+    best_peaks = np.sort(peaks[top_indices])
+    if len(best_peaks) < n:
+        filler = np.linspace(0, T-1, n)
+        best_peaks = filler # Fallback
+    best_start_times = torch.tensor(best_peaks, dtype=derivative_preds_np.dtype, device=device)
+
+    # --- Mock DP Matrix for Plotting ---
+    dp_mat = torch.full((n, T, T), float(-1e9), device=device)
+    for i, peak_time in enumerate(best_peaks):
+        dp_mat[i, :, int(peak_time)] = 1.0 
+
+    # --- Identical Plotting Logic ---
+    dp_mat_cpu = dp_mat.detach().cpu()
+    best_start_times_cpu = best_start_times.detach().cpu().numpy()
+    dp_to_plot = dp_mat_cpu.max(dim=1)[0].numpy()
+    
+    masked_dp = np.ma.masked_where(dp_to_plot <= -1e8, dp_to_plot)
+
+    plt.figure(figsize=(12, 6))
+    cmap = plt.cm.viridis
+    cmap.set_bad(color='white')
+    
+    plt.imshow(masked_dp, aspect='auto', origin='lower', cmap=cmap)
+    plt.colorbar(label='Peak Detection (Naive)')
+    plt.xlabel('End time (frame)')
+    plt.ylabel('Phoneme index')
+    plt.title('Naive Peak Detection (Ablation)')
+    plt.plot(best_start_times_cpu, range(len(best_start_times_cpu)), 'r.-', label='Detected Peaks')
+    plt.legend()
+    plt.tight_layout()
+    save_path = 'peak_detection_ablation.png'
+    plt.savefig(save_path)
+    plt.close()
+    
+    print(f"Ablation plot saved as {save_path}")
+
+    return best_start_times
+
+# ------------------------- phoneme alignment main ------------------------
+def phoneme_alignment(p_seq, w_phi, original_lengths, len_ratio, derivative_preds_np, probs_real):
+    gamma = 1e-20
+    T = int(original_lengths[0])
+    n = len(p_seq)
+    device = derivative_preds_np.device
+
+    if isinstance(probs_real, np.ndarray):
+        probs_real = torch.tensor(probs_real, device=device)
+    cumsum_probs = torch.cumsum(probs_real, dim=0)
+    
+    
+    
+    phoneme_mappings = {p.lower(): timit_to_leehon_map_MACRO.get(p.lower(), 'sil') if p.lower() not in timit_leehon_39_phonemes else p.lower() for p in p_seq}
+    derivatives = torch.cat([torch.tensor([0], device=derivative_preds_np.device), torch.diff(derivative_preds_np, dim=0)])
+
+    dp_mat = torch.full((n, T, T), float(-1e9), device=device)
+
+    p_idx0 = phoneme_to_idx_MACRO[phoneme_mappings[p_seq[0].lower()]]
+
+    # Vectorized init for first phoneme
+    t_e = torch.arange(T, device=device)
+    dp_mat[0, 0, :] = (
+        w_phi[0] * compute_phi_1(derivatives, 0, t_e)
+        + w_phi[1] * compute_phi_1(derivatives, 0, t_e)
+
+    )
+
+    for i in tqdm(range(1, n)):
+        p_idx = phoneme_to_idx_MACRO[phoneme_mappings[p_seq[i].lower()]]
+        # Vectorized over t_start and t_end
+        t_start = torch.arange(T, device=device)
+        t_end = torch.arange(T, device=device)
+        t_start_grid, t_end_grid = torch.meshgrid(t_start, t_end, indexing='ij')
+        valid_mask = t_start_grid < t_end_grid
+
+        phi1_dev = compute_phi_1(derivatives, t_start_grid, t_end_grid)
+        phi1 = compute_phi_1(derivative_preds_np, t_start_grid, t_end_grid)
+        phi2 = compute_phi_2(cumsum_probs, p_idx, t_start_grid, t_end_grid)
+        total_phi = w_phi[0] * phi1_dev + w_phi[1] * phi2
+
+        # Max over all possible previous end times.
+        # Vectorized equivalent of the original per-t_end loop: for each previous
+        # end time s, logsumexp over dp_mat[i-1]'s start-rows is the SAME regardless
+        # of t_end (invalid spans start>=end are -1e9 and never contribute), so the
+        # O(T) inner loop over t_end collapses to one column-wise logsumexp + mask.
+        # Bit-identical output; ~T x faster on long utterances.
+        col_lse = torch.logsumexp(dp_mat[i-1] / gamma, dim=0) * gamma  # (T,) over start rows
+        prev_scores = torch.where(
+            valid_mask,
+            col_lse.unsqueeze(1).expand(T, T),
+            torch.full((T, T), float(-1e9), device=device),
+        )
+        dp_mat[i] = torch.where(valid_mask, total_phi + prev_scores, torch.full_like(total_phi, float(-1e9)))
+
+    # Backtracking
+    best_start_times = torch.zeros((n), dtype=derivative_preds_np.dtype, device=device)
+    best_prev_t_end = T-1
+    for i in range(n):
+        cur_ph = n-1-i
+        scores = dp_mat[cur_ph, :, best_prev_t_end]
+        soft_weights = torch.softmax(scores / gamma, dim=0)
+        expected_idx = (soft_weights * torch.arange(T, device=device, dtype=derivative_preds_np.dtype)).sum()
+        best_start_times[cur_ph] = expected_idx      
+        best_prev_t_end = int(expected_idx.round().item())
+        
+    # DP-matrix figure is a debug artifact only; best_start_times (the return
+    # value) is already computed above. Skip entirely unless an output dir is set
+    # (the web demo sets it). Saves ~570ms/utterance in training/eval.
+    if _dp_matrix_out_dir is not None:
+        dp_mat_cpu = dp_mat.detach().cpu()
+        best_start_times_cpu = best_start_times.detach().cpu().numpy()
+        dp_to_plot = dp_mat_cpu.max(dim=1)[0].numpy()
+        masked_dp = np.ma.masked_where(dp_to_plot <= -1e8, dp_to_plot)  # mask all values <= -1e8
+
+        plt.figure(figsize=(12, 6))
+        cmap = plt.cm.viridis
+        cmap.set_bad(color='white')
+        real_min = masked_dp.min()
+        real_max = masked_dp.max()
+        # Plot DP matrix (max over start times)
+        plt.imshow(masked_dp, aspect='auto', origin='lower', cmap=cmap, vmin=real_min, vmax=real_max)
+        plt.colorbar(label='DP Score (max over start)')
+        plt.xlabel('End time (frame)')
+        plt.ylabel('Phoneme index')
+        plt.title('DP Matrix with Best Path')
+        # Overlay best_start_times as a red line
+        plt.plot(best_start_times_cpu, range(len(best_start_times_cpu)), 'r.-', label='Best start times')
+        plt.legend()
+        plt.tight_layout()
+        _save_path = os.path.join(_dp_matrix_out_dir or '.', 'dp_matrix_with_path.png')
+        plt.savefig(_save_path)
+        print(f"DP matrix with path plot saved as {_save_path}")
+        
+    return best_start_times
+
+def compute_phi_1(derivative_preds_np: torch.Tensor, t_start: Union[torch.Tensor, int], t_end: Union[torch.Tensor, int]) -> torch.Tensor:
+    """
+    Computes phi_1 for dynamic programming.
+    t_start and t_end can be scalars or tensors of the same shape.
+    Returns a tensor of scores.
+    """
+    # Ensure t_start and t_end are tensors
+    t_start = torch.as_tensor(t_start, device=derivative_preds_np.device)
+    t_end = torch.as_tensor(t_end, device=derivative_preds_np.device)
+    # Broadcast to same shape
+    t_start, t_end = torch.broadcast_tensors(t_start, t_end)
+    # Valid indices
+    valid = (t_end < derivative_preds_np.shape[0]-1) & (t_start < derivative_preds_np.shape[0]-1) & (t_end > 0) & (t_start > 0)
+    score = torch.zeros_like(t_start, dtype=derivative_preds_np.dtype, device=derivative_preds_np.device)
+
+    eps = 1e-6
+    tanh_scale = 1e-3 #1e-2 #0.5
+    if valid.any():
+        
+        # start_pos - 
+        idx_s = t_start[valid].long()
+        s_center = torch.tanh(tanh_scale * derivative_preds_np[idx_s])
+        s_prev = torch.tanh(tanh_scale * derivative_preds_np[idx_s -1])
+        s_next = torch.tanh(tanh_scale * derivative_preds_np[idx_s +1])
+        delta_prev_s = s_center - s_prev
+        delta_next_s = s_center - s_next
+        scores_zerocross_s = (1-torch.sqrt(s_center**2)) +  torch.sqrt(delta_prev_s **2 + eps) + torch.sqrt(delta_next_s**2 + eps)
+        # orig - 
+        score[valid] += scores_zerocross_s
+
+        # end_pos - 
+        idx_e = t_end[valid].long()
+        e_center = torch.tanh(tanh_scale * derivative_preds_np[idx_e]) #do i need this? not sure
+        e_prev = torch.tanh(tanh_scale * derivative_preds_np[idx_e -1])
+        e_next = torch.tanh(tanh_scale * derivative_preds_np[idx_e +1])
+        delta_prev_e = e_center - e_prev
+        delta_next_e = e_center - e_next
+        scores_zerocross_e = (1-torch.sqrt(e_center**2)) + torch.sqrt(delta_prev_e **2 + eps) + torch.sqrt(delta_next_e**2 + eps)
+        # orig - 
+        score[valid] += scores_zerocross_e
+    return score
+
+def compute_phi_2(cumsum_probs: torch.Tensor, p: int, t_start: Union[torch.Tensor, int], t_end: Union[torch.Tensor, int]) -> torch.Tensor:
+    """
+    Computes phi_2 for dynamic programming.
+    t_start and t_end can be scalars or tensors of the same shape.
+    Returns a tensor of scores.
+    """
+    t_start = torch.as_tensor(t_start, device=cumsum_probs.device)
+    t_end = torch.as_tensor(t_end, device=cumsum_probs.device)
+    t_start, t_end = torch.broadcast_tensors(t_start, t_end)
+    # Valid indices
+    valid = (t_end < cumsum_probs.shape[0]) & (t_start < cumsum_probs.shape[0]) & (t_end > 0) & (t_start >= 0)
+    probs_score = torch.zeros_like(t_start, dtype=cumsum_probs.dtype, device=cumsum_probs.device)
+    # Only assign where valid
+    probs_score[valid] = cumsum_probs[t_end[valid], p] - torch.where(
+        t_start[valid] > 0,
+        cumsum_probs[t_start[valid], p],
+        torch.zeros_like(t_start[valid], dtype=cumsum_probs.dtype, device=cumsum_probs.device)
+    )
+    lengths = (t_end - t_start).clamp(min=1)
+    probs_score[valid] = probs_score[valid] / lengths[valid]
+    return (probs_score)
+
+def best_phoneme_for_segments(cumsum_probs: torch.Tensor, t_start: torch.Tensor, t_end: torch.Tensor):
+    """
+    For each (t_start, t_end) pair (tensors broadcasted to same shape),
+    compute the average probability per phoneme over the segment and return:
+      - max_vals: tensor of shape (pairs,) with the max average prob per pair
+      - max_idx:  LongTensor of shape (pairs,) with the argmax phoneme index per pair
+    """
+    device = cumsum_probs.device
+    t_start = torch.as_tensor(t_start, device=device)
+    t_end = torch.as_tensor(t_end, device=device)
+    t_start, t_end = torch.broadcast_tensors(t_start, t_end)
+
+    valid = (t_end < cumsum_probs.shape[0]) & (t_start < cumsum_probs.shape[0]) & (t_end > 0) & (t_start >= 0)
+    max_vals = torch.zeros_like(t_start, dtype=cumsum_probs.dtype, device=device)
+    max_idx = torch.full_like(t_start, -1, dtype=torch.long, device=device)
+
+    if not valid.any():
+        return max_vals, max_idx
+
+    idx_end = t_end[valid].long()
+    idx_start = t_start[valid].long()
+    probs_end = cumsum_probs[idx_end]  # (k, P)
+    probs_start = torch.zeros_like(probs_end)
+    nonzero_mask = idx_start > 0
+    if nonzero_mask.any():
+        probs_start[nonzero_mask] = cumsum_probs[idx_start[nonzero_mask]]
+
+    segment_sum = probs_end - probs_start  # (k, P)
+    lengths = (t_end[valid] - t_start[valid]).clamp(min=1).unsqueeze(1).to(segment_sum.dtype)
+    segment_mean = segment_sum / lengths  # (k, P)
+
+    vals, idxs = segment_mean.max(dim=1)  # per-row max and argmax
+    max_vals[valid] = vals
+    max_idx[valid] = idxs.long()
+    return max_vals, max_idx
+
+def get_timit_61_phoneme_mappings():
+    """
+    Returns the TIMIT 61 phoneme-to-index mapping and the reverse index-to-phoneme mapping.
+
+    Returns:
+        phoneme_to_idx (dict): Dictionary mapping phonemes to unique indices.
+        idx_to_phoneme (dict): Dictionary mapping indices to their corresponding phonemes.
+    """
+    # this is actually including the leehon 39 phonemes!!!!!
+    timit_61_phonemes = [
+        'aa', 'ae', 'ah', 'ao', 'aw', 'ax', 'ax-h', 'axr', 'ay',
+        'b', 'bcl', 'ch', 'd', 'dcl', 'dh', 'dx', 'eh', 'el', 'em', 'en', 'eng', 'epi', 'er', 'ey',
+        'f', 'g', 'gcl', 'h#', 'hh', 'hv', 'ih', 'ix', 'iy', 'jh', 'k', 'kcl', 'l', 'm', 'n', 'ng',
+        'nx', 'ow', 'oy', 'p', 'pau', 'pcl', 'q', 'r', 's', 'sh', 't', 'tcl', 'th', 'uh', 'uw', 'ux',
+        'v', 'w', 'y', 'z', 'zh'
+    ]
+    timit_leehon_39_phonemes = [
+        'ao', 'ae', 'ah','aw', 'er', 'ay', 
+        'b', 'sil', 'ch', 'd', 'dh', 'dx', 'eh', 'el', 'm', 'en', 'ng', 'ey',
+        'f', 'g', 'hh', 'ih', 'iy', 'jh', 'k', 'v', 'w', 'y', 'z', 'sh', 't', 'r', 's', 'th','uh', 'uw', 'oy', 'ow','p'
+    ]
+    # Create mappings
+    phoneme_to_idx = {phoneme: idx for idx, phoneme in enumerate(timit_leehon_39_phonemes)}
+    idx_to_phoneme = {idx: phoneme for phoneme, idx in phoneme_to_idx.items()}
+
+    return phoneme_to_idx, idx_to_phoneme
+
+# --------------------------------
+
+
+def timit_to_leehon(timit_label):
+    # Mapping of TIMIT 61 phonemes to Leehon 39 phonemes
+    timit_to_leehon_map = {
+        'aa': 'ao', 'ae': 'ae', 'ah': 'ah', 'ao': 'ao', 'aw': 'aw', 'ax': 'ah', 'ax-h': 'ah', 'axr': 'er', 'ay': 'ay',
+        'b': 'b', 'bcl': 'sil', 'ch': 'ch', 'd': 'd', 'dcl': 'sil', 'dh': 'dh', 'dx': 'dx', 'eh': 'eh', 'el': 'el',
+        'em': 'm', 'en': 'en', 'eng': 'ng', 'epi': 'sil', 'er': 'er', 'ey': 'ey', 'f': 'f', 'g': 'g', 'gcl': 'sil',
+        'h#': 'sil', 'hh': 'hh', 'hv': 'hh', 'ih': 'ih', 'ix': 'ih', 'iy': 'iy', 'jh': 'jh', 'k': 'k', 'kcl': 'sil',
+        'l': 'el', 'm': 'm', 'n': 'en', 'ng': 'ng', 'nx': 'en', 'ow': 'ow', 'oy': 'oy', 'p': 'p', 'pau': 'sil', 'pcl': 'sil',
+        'q': 't', 'qcl': 'sil', 'r': 'r', 's': 's', 'sh': 'sh', 't': 't', 'tcl': 'sil', 'th': 'th', 'uh': 'uh', 'uw': 'uw',
+        'ux': 'uw', 'v': 'v', 'w': 'w', 'y': 'y', 'z': 'z', 'zh': 'sh', '':'sil'
+    }
+    
+    # Return the corresponding Leehon 39 label, or None if the label is not found
+    return timit_to_leehon_map.get(timit_label.lower(), None)
+
+def load_phoneme_stats():
+    phonemes_path = "phonemes_39"
+    stats_path = "phoneme_stats_39.out"
+    
+    # Load phoneme names
+    with open(phonemes_path, "r") as f:
+        phonemes = [line.strip() for line in f]
+
+    # Load mu values (second row of stats file)
+    with open(stats_path, "r") as f:
+        lines = f.readlines()
+        mu_values = list(map(float, lines[1].strip().split()))  # Convert to float
+        sigma_values = list(map(float, lines[2].strip().split()))  # Convert to float
+    
+    # Create phoneme-to-mu dictionary
+    phoneme_mu_dict = dict(zip(phonemes, mu_values))
+    phoneme_sigma_dict = dict(zip(phonemes, sigma_values))
+    
+    return phoneme_mu_dict, phoneme_sigma_dict
+
+# Load phoneme stats once
+
+
+def get_mu_stats(p):
+    phoneme_mu_dict, _ = load_phoneme_stats()
+    """Return the mu value for the given phoneme p."""
+    return phoneme_mu_dict.get(p, None)  # Return None if phoneme is not found
+
+def get_sigma_stats(p):
+    _, phoneme_sigma_dict = load_phoneme_stats()
+    """Return the mu value for the given phoneme p."""
+    return phoneme_sigma_dict.get(p, None)  # Return None if phoneme is not found
+
+def replicate_first_k_frames(x, k, dim):
+    return torch.cat([x.index_select(dim=dim, index=torch.LongTensor([0] * k).to(x.device)), x], dim=dim)
+
+class LambdaLayer(nn.Module):
+    def __init__(self, lambd):
+        super(LambdaLayer, self).__init__()
+        self.lambd = lambd
+    def forward(self, x):
+        return self.lambd(x)
+
+class PrintShapeLayer(nn.Module):
+    def __init__(self):
+        super(PrintShapeLayer, self).__init__()
+    def forward(self, x):
+        print(x.shape)
+        return x
+
+def length_to_mask(length, max_len=None, dtype=None):
+    """length: B.
+    return B x max_len.
+    If max_len is None, then max of length will be used.
+    """
+    assert len(length.shape) == 1, 'Length shape should be 1 dimensional.'
+    max_len = max_len or length.max().item()
+    mask = torch.arange(max_len, device=length.device,
+                        dtype=length.dtype).expand(len(length), max_len) < length.unsqueeze(1)
+    if dtype is not None:
+        mask = torch.as_tensor(mask, dtype=dtype, device=length.device)
+    return mask
+
+def detect_peaks_worker(xi,w_phi, p_seq, original_lengths, probs_real, len_ratio, width, distance):
+    print(f"num peaks = {len(p_seq)}")
+    print(f"xi type: {type(xi)}")
+    preds_np = xi.requires_grad_(True)
+    median_h = preds_np.median()
+    preds_np = preds_np - median_h
+    derivative_preds_np = preds_np
+    xmin, xmax = xi.min(), xi.max()
+    xi = (xi - xmin) / (xmax - xmin)
+    xi = xi.flatten()
+    
+    peaks = phoneme_alignment(p_seq,w_phi, original_lengths, len_ratio, derivative_preds_np, probs_real)
+    
+    if len(peaks) == 0:
+        peaks = torch.tensor([xi.shape[0] - 1], device=xi.device)
+    
+    return peaks
+
+def detect_peaks(x,w_phi, original_lengths_all, phonemes, len_ratio, probs_real_all):
+    """Detect peaks of next_frame_classifier using multithreading."""
+    
+    out = []
+    with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
+        xi=x
+        p_seq = phonemes
+        original_lengths = original_lengths_all
+        probs_real = probs_real_all
+        if len(xi)!=0:
+            result = detect_peaks_worker(xi, w_phi, p_seq, [original_lengths], probs_real, len_ratio, width=None, distance=None)
+        out.append(result)
+    
+    return out
+
+class PrecisionRecallMetric:
+    def __init__(self):
+        self.precision_counter = 0
+        self.recall_counter = 0
+        self.pred_counter = 0
+        self.gt_counter = 0
+        self.eps = 1e-5
+        self.data = []
+        self.tolerance = 2
+        self.width_range = [1]
+        self.distance_range = [1]
+
+    def get_metrics(self, precision_counter, recall_counter, pred_counter, gt_counter):
+        EPS = 1e-7
+        
+        precision = precision_counter / (pred_counter + self.eps)
+        recall = recall_counter / (gt_counter + self.eps)
+        f1 = 2 * (precision * recall) / (precision + recall + self.eps)
+        
+        os = recall / (precision + EPS) - 1
+        r1 = np.sqrt((1 - recall) ** 2 + os ** 2)
+        r2 = (-os + recall - 1) / (np.sqrt(2))
+        rval = 1 - (np.abs(r1) + np.abs(r2)) / 2
+
+        return precision, recall, f1, rval
+
+    def zero(self):
+        self.data = []
+        
+    def update(self, seg, pos_pred, length,original_lengths_all, probs_all,phonemes_all):
+        for seg_i, pos_pred_i, length_i , original_length, probs,phonemes in zip(seg, pos_pred, length,original_lengths_all,probs_all,phonemes_all):
+            self.data.append((seg_i, pos_pred_i, length_i.item(),[original_length.item()], probs, phonemes))
+
+
+    def get_stats(self, width=None, distance=None):
+        print(f"calculating metrics using {len(self.data)} entries")
+        max_rval = -float("inf")
+        min_l1_dist = float("inf")
+        best_params = None
+        segs = list(map(lambda x: x[0], self.data))
+        length = list(map(lambda x: x[2], self.data))
+        yhats = list(map(lambda x: x[1], self.data))
+        original_lengths_all = list(map(lambda x: x[3], self.data))
+        probs = list(map(lambda x: x[4], self.data))
+        phonemes = list(map(lambda x: x[5], self.data))
+
+        width_range = self.width_range
+        distance_range = self.distance_range
+
+        if width is not None:
+            width_range = [width]
+            distance_range = [distance]
+        sr = 16000
+        len_ratio = 161.34011627906978
+
+        for width in width_range:
+            for distance in distance_range:
+                    for (y, yhat,original_len, phoneme, prob) in zip(segs, yhats, original_lengths_all, phonemes, probs):
+                        if isinstance(y,list):
+                            y = torch.tensor(y, device=yhat.device, dtype=yhat.dtype)
+                        peaks = detect_peaks(x=yhat,w_phi= [0.5,0.5],
+                                             original_lengths_all = original_len[0],
+                                             phonemes = phoneme,
+                                             len_ratio = 161.34011627906978 , 
+                                             probs_real_all = prob)
+                        peaks = peaks[0]* len_ratio/sr
+                        yhat = peaks
+                        yhat = yhat[1:]
+
+                        if isinstance(y,list):
+                            y = torch.tensor(y, device=yhat.device, dtype=yhat.dtype)
+                            y = y*len_ratio/sr
+                        l1_dist = torch.mean(torch.abs(y - yhat)).item()
+                        l2_dist = torch.mean((y - yhat)**2).item()
+                        if l1_dist<min_l1_dist:
+                            min_l1_dist = l1_dist
+                            out = (l1_dist,l2_dist)
+                            best_params = width, distance
+        self.zero()
+        print(f"best peak detection params: {best_params} (width, distance)")
+        print(f"best peak detection L1_DIST: {l1_dist}")
+        print(f"best peak detection L2_DIST: {l2_dist}")
+        return out, best_params
+
+
+class StatsMeter:
+    def __init__(self):
+        self.data = []
+
+    def update(self, item):
+        if type(item) == list:
+            self.data.extend(item)
+        else:
+            self.data.append(item)
+
+    def get_stats(self):
+        data = np.array(self.data)
+        if len(data)==0:
+            return float('nan')
+        mean = data.mean()
+        return mean
+
+    def zero(self):
+        self.data.clear()
+        assert len(self.data) == 0, "StatsMeter didn't clear"
+
+
+class Timer:
+    def __init__(self, msg):
+        self.msg = msg
+        self.start_time = None
+
+    def __enter__(self):
+        self.start_time = time.time()
+        print(f"{self.msg} -- started")
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        print(f"{self.msg} -- done in {(time.time() - self.start_time)} secs")
+
+
+def max_min_norm(x):
+    x -= x.min(-1, keepdim=True)[0]
+    x /= x.max(-1, keepdim=True)[0]
+    return x
+
+
+def line():
+    print(90 * "-")

word_g2p.py

@@ -0,0 +1,125 @@
+"""
+Stand-alone word -> Lee-Hon-39 phoneme front-end for *word-level* alignment.
+
+This converts orthographic words into a phoneme sequence so the (otherwise
+phoneme-level) aligner can run on word input. It is deliberately independent
+of MFA / any user-supplied pronunciation dictionary or acoustic model:
+
+    word  --espeak (grapheme->IPA, rule-based, multilingual)-->  IPA
+    IPA   --panphon ipa_segs-->  IPA segments
+    seg   --dutch_preprocess.find_best_leehon39 (articulatory distance)-->  LH39
+
+The IPA->LH39 step is the *same* panphon mapping already used on the phoneme
+path, so word-level and phoneme-level share one LH39 back-end. Only the word
+case routes through here; the phoneme path is unchanged.
+
+espeak is a stand-alone open-source phonemizer (not MFA), so nothing here
+depends on the system we compare against at runtime.
+"""
+import os
+import subprocess
+
+import panphon
+
+import dutch_preprocess
+
+# espeak voice used for word-level G2P. "en-us" = American English (TIMIT).
+# Override per-language via the FDNFA_G2P_VOICE env var (e.g. "nl", "de", "he").
+DEFAULT_VOICE = os.environ.get("FDNFA_G2P_VOICE", "en-us")
+
+# Insert TIMIT-style closure/silence segments so the word-derived phoneme
+# sequence matches the model's training granularity. Closures/silences are ~22%
+# of TIMIT reference segments (all map to LH39 'sil') and no text G2P emits
+# them; restoring them via this parameter-free rule recovers most of the
+# word-vs-phoneme gap (+~16 pts @25ms). Disable with FDNFA_WORD_CLOSURES=0.
+USE_CLOSURES = os.environ.get("FDNFA_WORD_CLOSURES", "1").lower() not in ("0", "false", "no")
+STOPS = {"b", "d", "g", "p", "t", "k"}
+
+_ft = panphon.FeatureTable()
+_cache = {}  # (word_lower, voice) -> [lh39, ...]
+# espeak decorates IPA with stress / length / syllable marks that are not
+# phonemes; strip them before segmentation.
+_STRIP = dict.fromkeys(map(ord, "ˈˌːˑ.‿|"), None)
+
+
+def _with_closures(seq):
+    """Insert a 'sil' (closure) before every stop, mirroring TIMIT segmentation."""
+    out = []
+    for p in seq:
+        if p in STOPS:
+            out.append("sil")
+        out.append(p)
+    return out
+
+
+def _espeak_ipa(word, voice):
+    try:
+        out = subprocess.run(
+            ["espeak", "-q", "--ipa", "-v", voice, word],
+            capture_output=True, text=True, timeout=10,
+        ).stdout
+    except Exception as exc:  # espeak missing / failed -> caller handles empty
+        print(f"[word_g2p] espeak failed for {word!r}: {exc}")
+        return ""
+    return out.strip().translate(_STRIP)
+
+
+# Word-level G2P backend: "espeak" (default, multilingual, MFA-independent) or
+# "mfa" (the english_us_arpa Pynini G2P that Montreal Forced Aligner uses — for
+# apples-to-apples English word-level comparison). Override via FDNFA_WORD_G2P.
+G2P_BACKEND = os.environ.get("FDNFA_WORD_G2P", "espeak").lower()
+
+
+def word_to_lh39(word, voice=None, backend=None):
+    """Orthographic word -> list of LH39 phonemes (cached).
+
+    Routes through the espeak or MFA-english_us_arpa G2P. `backend`/`voice`
+    default to the FDNFA_WORD_G2P / FDNFA_G2P_VOICE env vars when not passed, so a
+    long-running process (e.g. the Gradio app) can switch per call by argument.
+    The MFA backend applies to English (it uses the english_us_arpa dictionary).
+    """
+    if backend is None:
+        backend = os.environ.get("FDNFA_WORD_G2P", "espeak").lower()
+    if voice is None:
+        voice = os.environ.get("FDNFA_G2P_VOICE", DEFAULT_VOICE)
+    if backend == "mfa":
+        import mfa_g2p  # lazy: mfa_g2p imports closure helpers from this module
+        return mfa_g2p.word_to_lh39_mfa(word, voice=voice)
+    if backend == "char":
+        return _char_word_lh39(word)
+    return _espeak_word_lh39(word, voice)
+
+
+def _char_word_lh39(word):
+    """Word -> LH39 with NO G2P model: segment the (romanized) characters with
+    panphon and map each directly to LH39 by articulatory-feature distance. Used
+    for languages with no MFA model and where a grapheme-to-phoneme converter is
+    deliberately avoided (e.g. Hebrew romanized transcripts)."""
+    key = (word.lower(), "char")
+    if key in _cache:
+        return _cache[key]
+    segs = _ft.ipa_segs(word.lower())
+    lh39 = [dutch_preprocess.find_best_leehon39(s)[0] for s in segs if s.strip()]
+    if not lh39:
+        lh39 = ["sil"]
+    if USE_CLOSURES:
+        lh39 = _with_closures(lh39)
+    _cache[key] = lh39
+    return lh39
+
+
+def _espeak_word_lh39(word, voice=DEFAULT_VOICE):
+    """Word -> LH39 via espeak (cached). The espeak branch of word_to_lh39, also
+    used by the MFA backend as the OOV fallback for languages with no MFA G2P."""
+    key = (word.lower(), voice)
+    if key in _cache:
+        return _cache[key]
+    ipa = _espeak_ipa(word, voice)
+    segs = _ft.ipa_segs(ipa) if ipa else []
+    lh39 = [dutch_preprocess.find_best_leehon39(s)[0] for s in segs if s.strip()]
+    if not lh39:
+        lh39 = ["sil"]
+    if USE_CLOSURES:
+        lh39 = _with_closures(lh39)
+    _cache[key] = lh39
+    return lh39