JointRecAlign 0.1.0__py3-none-any.whl

jointrecalign/__init__.py

@@ -0,0 +1 @@
+"""Joint phone recognition and alignment with wav2vec2 (CTC + cross-attention aligner)."""

jointrecalign/__main__.py

@@ -0,0 +1,4 @@
+"""`python -m jointrecalign`: same as the JointRecAlign command."""
+from .transcribe_cli import main
+
+main()

jointrecalign/download.py

@@ -0,0 +1,82 @@
+"""The released model: downloaded once, on first use, into a cache folder.
+
+    Linux    ~/.cache/JointRecAlign
+    macOS    ~/Library/Caches/JointRecAlign
+    Windows  %LOCALAPPDATA%\\JointRecAlign
+
+JOINTRECALIGN_HOME overrides the folder; JOINTRECALIGN_MODEL_URL the download URL.
+"""
+import hashlib
+import os
+import shutil
+import sys
+import urllib.request
+import zipfile
+
+MODEL_NAME = "aligner_10ms_sil_alpha0.1"
+MODEL_URL = ("https://github.com/Imenbaa/Joint_phoneme_recognition_alignment/releases/download/"
+             f"v0.1.0/{MODEL_NAME}.zip")
+MODEL_SHA256 = "a1afc8f77544475e062b9de252abfe2d1b9f10c79f5b9f4e702c57548ee67f7a"
+
+
+def cache_dir():
+    if os.environ.get("JOINTRECALIGN_HOME"):
+        return os.environ["JOINTRECALIGN_HOME"]
+    if sys.platform == "win32":
+        base = os.environ.get("LOCALAPPDATA", os.path.expanduser("~"))
+    elif sys.platform == "darwin":
+        base = os.path.expanduser("~/Library/Caches")
+    else:
+        base = os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache"))
+    return os.path.join(base, "JointRecAlign")
+
+
+def _fetch(url, dest, log):
+    """Download url to dest, printing the progress every 5 %."""
+    with urllib.request.urlopen(url) as r, open(dest, "wb") as f:
+        total = int(r.headers.get("Content-Length") or 0)
+        done, shown = 0, -5
+        while True:
+            block = r.read(1 << 20)
+            if not block:
+                break
+            f.write(block)
+            done += len(block)
+            if total and 100 * done // total >= shown + 5:
+                shown = 100 * done // total
+                log(f"  {shown:3d} %  ({done / 1e6:.0f} / {total / 1e6:.0f} MB)")
+
+
+def default_model(log=print):
+    """Folder of the released model; downloaded, checked and unzipped the first time."""
+    root = cache_dir()
+    folder = os.path.join(root, MODEL_NAME)
+    if os.path.exists(os.path.join(folder, "model.safetensors")):
+        return folder
+    os.makedirs(root, exist_ok=True)
+    url = os.environ.get("JOINTRECALIGN_MODEL_URL", MODEL_URL)
+    part = os.path.join(root, MODEL_NAME + ".zip.part")
+    log(f"First use: downloading the model (1.2 GB, only once) into {root}")
+    try:
+        _fetch(url, part, log)
+    except Exception as e:
+        raise RuntimeError(f"could not download the model from {url} ({e}).\n"
+                           "Check the internet connection and try again, or download the zip by hand, "
+                           "unzip it and give its folder with --model.") from None
+    h = hashlib.sha256()
+    with open(part, "rb") as f:
+        for block in iter(lambda: f.read(1 << 20), b""):
+            h.update(block)
+    if h.hexdigest() != MODEL_SHA256 and "JOINTRECALIGN_MODEL_URL" not in os.environ:
+        os.remove(part)
+        raise RuntimeError("the downloaded model is damaged (checksum mismatch); please try again.")
+    log("Unzipping the model")
+    tmp = folder + ".tmp"
+    shutil.rmtree(tmp, ignore_errors=True)
+    with zipfile.ZipFile(part) as z:
+        z.extractall(tmp)
+    shutil.rmtree(folder, ignore_errors=True)
+    os.replace(os.path.join(tmp, MODEL_NAME), folder)
+    shutil.rmtree(tmp, ignore_errors=True)
+    os.remove(part)
+    return folder

jointrecalign/evaluation/__init__.py

@@ -0,0 +1 @@
+"""Alignment evaluation: metrics, attention/DTW, latent-SIL Viterbi and CTC forced-alignment decoders."""

jointrecalign/evaluation/attention_eval.py

@@ -0,0 +1,214 @@
+#!/usr/bin/env python3
+"""Score a checkpoint on an evaluation manifest: DTW on the cross-attention energy.
+
+The model aligns its own CTC greedy decode (no reference phones, no prior),
+chunk by chunk: no padding (the model sees exactly [start, end]), onset-frame
+extraction with contiguous ends, no silence trimming. ctcfa_eval.py uses the
+same conventions, so the two are directly comparable.
+
+Manifest (JSON):
+    {"files": {
+        "<file_id>": {
+            "wav": "/path/to/audio.wav",                      # any rate, mono or not
+            "chunks": [{"start": 0.0, "end": 12.3}, ...],     # seconds, scored segments
+            "ref_intervals": [{"phoneme": "a", "start": 0.12, "end": 0.20}, ...],
+            "style": "read",                                  # optional grouping key
+            "pause_intervals": [[s, e], ...],                 # optional, SIL scoring
+            "excluded_intervals": [[s, e], ...]               # optional, SIL scoring
+        }, ...}}
+
+Frame timing is read from the model config (frame stride and receptive field),
+so the same code scores 20 ms and 10 ms models.
+"""
+import numpy as np, torch, soundfile as sf, soxr, unicodedata
+from itertools import groupby
+
+SR = 16000
+FRAME_STRIDE_S = 0.02          # wav2vec2: 320 samples @ 16 kHz (default only)
+
+
+def model_sr(model):
+    """Sampling rate the model was trained on (config.input_sampling_rate,
+    16 kHz when absent)."""
+    return int(getattr(model.config, "input_sampling_rate", SR) or SR)
+
+
+def frame_stride_s(model):
+    """Seconds per encoder frame: conv hop (samples) / input rate. 0.02 for the
+    standard model; 0.01 with last conv stride 1 or 32 kHz input.
+    Frame i starts at i * stride (checked: 10 ms frame 2i == 20 ms frame i)."""
+    return getattr(model.config, "inputs_to_logits_ratio", 320) / model_sr(model)
+
+
+def frame_time_offset_s(model):
+    """Fixed, architecture-only correction from frame index to boundary time.
+
+    Frame i of the conv front-end sees samples [i*hop, i*hop + RF) (no padding),
+    so its centre is at i*hop + RF/2. A label change at frame i is placed halfway
+    between the centres of frames i-1 and i:
+        boundary = i*hop + (RF - hop) / 2
+    RF and hop come from config.conv_kernel / conv_stride. Nothing is tuned on data:
+    +2.5 ms for the 20 ms model, +7.5 ms for 10 ms (last stride 1), +1.25 ms at 32 kHz.
+    """
+    rf, hop = 1, 1
+    for k, s in zip(model.config.conv_kernel, model.config.conv_stride):
+        rf += (k - 1) * hop
+        hop *= s
+    return (rf - hop) / 2 / model_sr(model)
+
+
+def resample(a, orig_sr, target_sr):
+    """As librosa.resample (default soxr_hq, length ceil(n * ratio)), without librosa."""
+    n = int(np.ceil(a.shape[-1] * float(target_sr) / orig_sr))
+    y = soxr.resample(a, orig_sr, target_sr, quality="soxr_hq")
+    y = y[:n] if len(y) >= n else np.pad(y, (0, n - len(y)))
+    return np.asarray(y, dtype=a.dtype)
+
+
+def to_model_rate(seg, model):
+    """16 kHz chunk -> the model's input rate. Audio is always loaded at 16 kHz
+    first, so a 32 kHz model sees the same 8 kHz band limit as its (upsampled
+    16 kHz) training data."""
+    sr = model_sr(model)
+    if sr == SR:
+        return seg
+    return resample(seg, SR, sr).astype(np.float32)
+
+
+# Hypothesis tokens rewritten before scoring (affricates of the training phone
+# set that the reference inventory writes as one fricative).
+HYP_FIX = {"ts": "s", "tʃ": "ʃ"}
+SKIP_TOK = ("", " ", "[", "|", "[PAD]", "[UNK]", "[SIL]")
+
+
+def load_audio(path):
+    a, sr = sf.read(path)
+    if a.ndim > 1:
+        a = np.mean(a, axis=1)
+    if sr != SR:
+        a = resample(a, sr, SR)
+    return a.astype(np.float32)
+
+
+def dtw_path(att):
+    """[T, P] -> [T] phone index per frame.
+
+    Steps [1,1] and [1,0]: no skips, every frame consumed once, every phone
+    gets at least one frame.
+    """
+    from librosa.sequence import dtw          # evaluation only (pip install ".[eval]")
+    T, P = att.shape
+    if P == 0 or T < P:
+        return None
+    if P == 1:
+        return np.zeros(T, dtype=np.int32)
+    _, wp = dtw(C=-att, step_sizes_sigma=np.array([[1, 1], [1, 0]]))
+    path = np.full(T, -1, dtype=np.int32)
+    for fr, ph in wp:                      # librosa returns the path reversed
+        if 0 <= fr < T and path[fr] < ph:
+            path[fr] = ph
+    last = 0
+    for t in range(T):
+        if path[t] < 0:
+            path[t] = last
+        else:
+            last = path[t]
+    return path
+
+
+@torch.inference_mode()
+def run_manifest(model, feature_extractor, tokenizer, manifest,
+                 device="cuda", max_files=None):
+    """-> {file_id: {ref_intervals, hyp_intervals, ref_seq, hyp_seq, style}}"""
+    was_training = model.training
+    model.eval()
+    store = {}
+    stride = frame_stride_s(model)
+    offset = frame_time_offset_s(model)   # receptive-field centre rule
+    try:
+        for nm, meta in list(manifest["files"].items())[:max_files]:
+            wav = torch.from_numpy(load_audio(meta["wav"])).float()
+            hyp = []
+            for ch in meta["chunks"]:
+                s0, s1 = int(ch["start"] * SR), int(ch["end"] * SR)
+                seg = wav[s0:s1]
+                if seg.numel() < 400:      # under one conv receptive field
+                    continue
+                # the extractor only normalises; pass its own rate so a 16 kHz
+                # extractor also serves a 32 kHz model
+                x = feature_extractor(to_model_rate(seg.numpy(), model),
+                                      sampling_rate=feature_extractor.sampling_rate,
+                                      return_tensors="pt", return_attention_mask=True)
+                o = model(input_values=x.input_values.to(device),
+                          attention_mask=x.attention_mask.to(device),
+                          labels=None,          # target = CTC greedy decode
+                          prior_weight=0.0,     # no beta-binomial prior
+                          ctc_weight=0.0, align_weight=0.0, contrastive_weight=0.0)
+                T, P = int(o.input_lengths[0]), int(o.text_lengths[0])
+                if P == 0:
+                    continue
+                # log_softmax then exp: softmax-then-log clips the losers to a
+                # constant floor and destroys their ordering.
+                att = np.exp(o.energy[0, :T, :P].float().log_softmax(-1).cpu().numpy())
+                path = dtw_path(att)
+                if path is None:
+                    continue
+                tokens = tokenizer.convert_ids_to_tokens(
+                    o.predicted_phone_ids[0, :P].cpu().tolist())
+
+                # run-length encode; keep each phone's ONSET frame
+                spans, n = [], 0
+                for idx, grp in groupby(path.tolist()):
+                    spans.append((idx, n))
+                    n += sum(1 for _ in grp)
+
+                ca = []
+                for pi, sframe in spans:
+                    ph = tokens[pi]
+                    if ph in SKIP_TOK:
+                        continue
+                    if len(ph) == 1 and unicodedata.combining(ph):
+                        if ca:
+                            ca[-1]["phoneme"] += ph
+                        continue
+                    ca.append({"phoneme": HYP_FIX.get(ph, ph),
+                               "start": s0 / SR + sframe * stride + offset,
+                               "end": None})
+                # contiguous ends: end_i = start_{i+1}; last = chunk end
+                for i in range(len(ca) - 1):
+                    ca[i]["end"] = ca[i + 1]["start"]
+                if ca:
+                    ca[-1]["end"] = s0 / SR + n * stride + offset
+                hyp.extend(ca)
+
+            hyp.sort(key=lambda x: x["start"])
+            store[nm] = {"ref_intervals": meta["ref_intervals"],
+                         "hyp_intervals": hyp,
+                         "ref_seq": [x["phoneme"] for x in meta["ref_intervals"]],
+                         "hyp_seq": [x["phoneme"] for x in hyp],
+                         "style": meta.get("style", "ALL")}
+    finally:
+        model.train(was_training)
+    return store
+
+
+def duration_fidelity(store, align_sequences, gate=0.5):
+    """Slope of hypothesis duration on reference duration, over matched phones.
+    1.0 = faithful; 0.0 = every phone gets the same length whatever it is.
+    The gate discards pairs the edit-distance matched across seconds of audio.
+    """
+    R, H = [], []
+    for d in store.values():
+        for i, j in align_sequences(d["ref_seq"], d["hyp_seq"]):
+            if i is None or j is None or d["ref_seq"][i] != d["hyp_seq"][j]:
+                continue
+            r, h = d["ref_intervals"][i], d["hyp_intervals"][j]
+            if abs(h["start"] - r["start"]) > gate or abs(h["end"] - r["end"]) > gate:
+                continue
+            R.append((r["end"] - r["start"]) * 1000)
+            H.append((h["end"] - h["start"]) * 1000)
+    if len(R) < 100:
+        return {}
+    R, H = np.array(R), np.array(H)
+    return {"dur_slope": float(np.polyfit(R, H, 1)[0]),
+            "dur_corr": float(np.corrcoef(R, H)[0, 1])}

jointrecalign/evaluation/ctcfa_eval.py

@@ -0,0 +1,131 @@
+"""Score a checkpoint with CTC forced alignment on its CTC head (no attention).
+
+Same manifest, chunks and output format as attention_eval.run_manifest, so the
+two are directly comparable; only the boundary extraction differs:
+
+  attention (run_manifest): DTW on the cross-attention energy, onset frames,
+                            contiguous ends.
+  ctcfa (this file):        greedy CTC decode, torchaudio forced_align of that
+                            sequence on the CTC log posteriors, one interval per
+                            non-blank run with its own start AND end frame
+                            (blank frames between phones stay unassigned).
+
+The decoded phone sequence is the same CTC greedy decode the attention path
+aligns, so PER matches the attention eval; only the boundaries change.
+
+end_mode="contiguous" instead ends each phone where the next one starts and
+the last phone of a chunk at the chunk end, exactly as run_manifest does for
+the DTW path, so blank frames go to the preceding phone.
+
+Frame timing follows attention_eval: frame stride and receptive-field offset
+are read from the model config, so frame f of a chunk starting at cs is at
+cs + f * stride + offset, for 20 ms and 10 ms models alike.
+"""
+import unicodedata
+
+import numpy as np
+import torch
+import torchaudio
+
+from .attention_eval import (SR, HYP_FIX, SKIP_TOK, load_audio, frame_stride_s,
+                             frame_time_offset_s, to_model_rate)
+
+
+def ctc_logits(o):
+    for name in ("logits", "ctc_logits"):
+        v = getattr(o, name, None)
+        if v is not None:
+            return v
+    raise AttributeError("no CTC logits in the model output")
+
+
+@torch.inference_mode()
+def run_manifest_ctcfa(model, feature_extractor, tokenizer, manifest,
+                       device="cuda", max_files=None, end_mode="span"):
+    """-> {file_id: {ref_intervals, hyp_intervals, ref_seq, hyp_seq, style}}"""
+    was_training = model.training
+    model.eval()
+    blank = tokenizer.pad_token_id
+    stride = frame_stride_s(model)
+    offset = frame_time_offset_s(model)
+    store, n_fail = {}, 0
+    try:
+        for nm, meta in list(manifest["files"].items())[:max_files]:
+            wav = torch.from_numpy(load_audio(meta["wav"])).float()
+            hyp = []
+            for ch in meta["chunks"]:
+                s0, s1 = int(ch["start"] * SR), int(ch["end"] * SR)
+                seg = wav[s0:s1]
+                if seg.numel() < 400:      # under one conv receptive field
+                    continue
+                x = feature_extractor(to_model_rate(seg.numpy(), model),
+                                      sampling_rate=feature_extractor.sampling_rate,
+                                      return_tensors="pt", return_attention_mask=True)
+                o = model(input_values=x.input_values.to(device),
+                          attention_mask=x.attention_mask.to(device),
+                          labels=None, prior_weight=0.0,
+                          ctc_weight=0.0, align_weight=0.0, contrastive_weight=0.0)
+                T = int(o.input_lengths[0])
+                logp = ctc_logits(o)[0, :T].float().log_softmax(-1).cpu()
+
+                # greedy collapse -> forced-alignment target
+                ids, prev = [], None
+                for t in logp.argmax(-1).tolist():
+                    if t != prev and t != blank:
+                        ids.append(t)
+                    prev = t
+                if not ids:
+                    continue
+                try:
+                    aligned, _ = torchaudio.functional.forced_align(
+                        logp.unsqueeze(0).contiguous(), torch.tensor([ids], dtype=torch.int32),
+                        torch.tensor([T], dtype=torch.int32), torch.tensor([len(ids)], dtype=torch.int32),
+                        blank=blank)
+                except Exception as e:     # e.g. more targets than frames
+                    n_fail += 1
+                    print(f"  [forced_align skipped] {nm} {ch['start']:.2f}-{ch['end']:.2f}: {e}")
+                    continue
+                aligned = aligned[0].tolist()
+
+                # runs with start AND end frame; blank frames stay unassigned
+                spans, run_id, run_start = [], None, 0
+                for fi, tid in enumerate(aligned):
+                    if tid != run_id:
+                        if run_id is not None and run_id != blank:
+                            spans.append((run_id, run_start, fi))
+                        run_id, run_start = tid, fi
+                if run_id is not None and run_id != blank:
+                    spans.append((run_id, run_start, len(aligned)))
+
+                off = s0 / SR + offset
+                ca = []
+                for tid, sf, ef in spans:
+                    ph = tokenizer.convert_ids_to_tokens(tid)
+                    if ph in SKIP_TOK:
+                        continue
+                    if len(ph) == 1 and unicodedata.combining(ph):
+                        if ca:
+                            ca[-1]["phoneme"] += ph
+                            ca[-1]["end"] = off + ef * stride
+                        continue
+                    ca.append({"phoneme": HYP_FIX.get(ph, ph),
+                               "start": off + sf * stride,
+                               "end": off + ef * stride})
+                if end_mode == "contiguous":
+                    for i in range(len(ca) - 1):
+                        ca[i]["end"] = ca[i + 1]["start"]
+                    if ca:
+                        ca[-1]["end"] = off + T * stride
+                hyp.extend(ca)
+
+            hyp.sort(key=lambda x: x["start"])
+            store[nm] = {"ref_intervals": meta["ref_intervals"],
+                         "hyp_intervals": hyp,
+                         "ref_seq": [x["phoneme"] for x in meta["ref_intervals"]],
+                         "hyp_seq": [x["phoneme"] for x in hyp],
+                         "style": meta.get("style", "ALL")}
+    finally:
+        model.train(was_training)
+    if n_fail:
+        print(f"  [ctcfa] {n_fail} chunks skipped by forced_align")
+    return store