| """Phase 2 shared tools: romanisation for MMS forced alignment, chunked aligner, pause normalisation, |
| pause / F0 / Yoruba-tone metrics. Used by prep2.py (data) and eval2.py (metrics).""" |
| import re, unicodedata |
| import numpy as np |
|
|
| |
| _MAP = {"ọ": "o", "ẹ": "e", "ṣ": "s", "ị": "i", "ụ": "u", "ṅ": "n", "ń": "n", "ǹ": "n", "ḿ": "m", "ɓ": "b", "ɗ": "d", |
| "ƙ": "k", "ƴ": "y", "ŋ": "n", "ɛ": "e", "ɔ": "o", "ß": "s", "æ": "ae", "œ": "oe", "ø": "o", "đ": "d", "ł": "l"} |
| TONE_MARK = {"\u0301": 2, "\u0300": 0, "\u0304": 1} |
| SENT_END = set(".?!") |
| CLAUSE = set(";:") |
| COMMA = set(",") |
| _VOW = set("aeiou") |
|
|
|
|
| def roman_char(c): |
| """one NFC char -> (romanised latin string, tone or None).""" |
| lc = c.lower() |
| if lc in _MAP: |
| return _MAP[lc], None |
| d = unicodedata.normalize("NFD", lc) |
| base = "".join(x for x in d if unicodedata.category(x) != "Mn") |
| tone = None |
| for x in d: |
| if x in TONE_MARK: |
| tone = TONE_MARK[x] |
| base = "".join(_MAP.get(b, b) for b in base) |
| base = "".join(b for b in base if ("a" <= b <= "z") or b == "'") |
| return base, tone |
|
|
|
|
| def words_for_align(text): |
| """text -> list of dict(word, rom, punct_after, chars=[(rom_char, orig_index, tone)]). |
| Words that romanise to nothing (digits, symbols) are aligned with the '*' star token.""" |
| out = [] |
| for m in re.finditer(r"\S+", text): |
| w = m.group(0) |
| chars = [] |
| for k, c in enumerate(w): |
| r, t = roman_char(c) |
| for j, rc in enumerate(r): |
| chars.append((rc, m.start() + k, t if j == len(r) - 1 else None)) |
| trail = re.sub(r"[\w'’]+", "", w[-3:]) if w else "" |
| p = None |
| if any(c in SENT_END for c in trail): |
| p = "sent" |
| elif any(c in CLAUSE for c in trail): |
| p = "clause" |
| elif any(c in COMMA for c in trail): |
| p = "comma" |
| rom = "".join(c[0] for c in chars) |
| out.append(dict(word=w, rom=rom if rom else "*", punct=p, chars=chars)) |
| return out |
|
|
|
|
| |
| class Aligner: |
| """torchaudio MMS_FA (CC-BY-NC weights, used only as a data tool/judge) with chunked emissions for long audio.""" |
|
|
| def __init__(self, device="cuda"): |
| import torch, torchaudio |
| self.torch = torch |
| self.F = torchaudio.functional |
| b = torchaudio.pipelines.MMS_FA |
| self.model = b.get_model().to(device).eval().half() |
| self.dict = b.get_dict() |
| self.dev = device |
| self.rate = 320 |
|
|
| def emission(self, y16, chunk_s=30, ctx_s=2): |
| torch = self.torch |
| sr = 16000 |
| x = torch.from_numpy(y16).float().to(self.dev) |
| if len(y16) <= (chunk_s + 2 * ctx_s) * sr: |
| with torch.inference_mode(): |
| e, _ = self.model(x[None].half()) |
| return e[0].float().log_softmax(-1) |
| outs = [] |
| step, ctx = chunk_s * sr, ctx_s * sr |
| nfr_total = len(y16) // self.rate |
| for s in range(0, len(y16), step): |
| a, b = max(0, s - ctx), min(len(y16), s + step + ctx) |
| with torch.inference_mode(): |
| e, _ = self.model(x[a:b][None].half()) |
| e = e[0].float().log_softmax(-1) |
| f0 = (s - a) // self.rate |
| nfr = min(step, len(y16) - s) // self.rate |
| outs.append(e[f0:f0 + nfr]) |
| e = torch.cat(outs, 0) |
| return e[:max(nfr_total, 1)] |
|
|
| def align(self, y16, words): |
| """-> per word (start_s, end_s, score) and per char spans [(start_s, end_s)], or None on failure.""" |
| torch = self.torch |
| em = self.emission(y16) |
| toks, wid = [], [] |
| for i, w in enumerate(words): |
| for c in w["rom"]: |
| if c in self.dict: |
| toks.append(self.dict[c]); wid.append(i) |
| if not toks or len(toks) * 1 > em.shape[0]: |
| return None |
| t = torch.tensor([toks], dtype=torch.int32, device=self.dev) |
| try: |
| ali, sc = self.F.forced_align(em[None], t, blank=0) |
| except Exception: |
| return None |
| ali, sc = ali[0].cpu().numpy(), sc[0].exp().cpu().numpy() |
| |
| spans, i, k = [], 0, 0 |
| T = len(ali) |
| while i < T: |
| if ali[i] == 0: |
| i += 1; continue |
| j = i |
| while j + 1 < T and ali[j + 1] == ali[i] and not (j + 1 < T and False): |
| j += 1 |
| spans.append((i, j + 1, float(sc[i:j + 1].mean()))) |
| i = j + 1 |
| |
| |
| if len(spans) != len(toks): |
| return None |
| r = self.rate / 16000 |
| W = [[None, None, []] for _ in words] |
| chars = [[] for _ in words] |
| for (a, b, s), w in zip(spans, wid): |
| if W[w][0] is None: |
| W[w][0] = a * r |
| W[w][1] = b * r; W[w][2].append(s); chars[w].append((a * r, b * r)) |
| res = [] |
| for i, (a, b, s) in enumerate(W): |
| if a is None: |
| return None |
| res.append(dict(start=a, end=b, score=float(np.mean(s)), chars=chars[i])) |
| return res |
|
|
|
|
| |
| def frame_db(y, sr, hop_s=0.01): |
| hop = int(sr * hop_s) |
| n = max(1, len(y) // hop) |
| fr = y[: n * hop].reshape(n, hop) |
| e = 10 * np.log10((fr.astype(np.float64) ** 2).mean(1) + 1e-10) |
| return e |
|
|
|
|
| def silence_mask(y, sr, rel_db=-40.0, hop_s=0.01): |
| e = frame_db(y, sr, hop_s) |
| ref = np.percentile(e, 98) |
| return e < ref + rel_db |
|
|
|
|
| def runs(mask): |
| """-> list of (start, end) index runs where mask is True.""" |
| out, i, n = [], 0, len(mask) |
| while i < n: |
| if mask[i]: |
| j = i |
| while j < n and mask[j]: |
| j += 1 |
| out.append((i, j)); i = j |
| else: |
| i += 1 |
| return out |
|
|
|
|
| PAUSE_TARGET = { |
| "sent": (0.30, 0.60), |
| "clause": (0.20, 0.50), |
| "comma": (0.00, 0.40), |
| None: (0.00, 0.20), |
| } |
| |
| |
| import os as _os |
| if _os.environ.get("P2PAUSE") == "B": |
| PAUSE_TARGET = {"sent": (0.35, 0.70), "clause": (0.25, 0.55), "comma": (0.00, 0.50), None: (0.00, 0.35)} |
| LEAD_S, TAIL_S = 0.12, 0.18 |
|
|
|
|
| def normalise_pauses(y, sr, words, ali, seg_a, seg_b, sil=None, lo_s=0.0, hi_s=None): |
| """Rebuild one segment [word seg_a .. seg_b] of y: edge silence fixed to LEAD_S/TAIL_S, every inter-word |
| silence clamped to PAUSE_TARGET[punct] (only truly silent frames are removed; silence is inserted, as a |
| copy of the quietest local frames, when a sentence/clause boundary has too little). -> (y_out, pauses)""" |
| hop = int(0.01 * sr) |
| if sil is None: |
| sil = silence_mask(y, sr) |
| hi_s = len(y) / sr if hi_s is None else hi_s |
| xf = int(0.005 * sr) |
| pieces, pauses = [], [] |
| a0 = ali[seg_a]["start"] |
| start = int(max(lo_s, a0 - LEAD_S, 0.0) * sr) |
| cur = start |
| for i in range(seg_a, seg_b): |
| e_i, s_n = ali[i]["end"], ali[i + 1]["start"] |
| f_a, f_b = int(e_i / 0.01), int(s_n / 0.01) + 1 |
| |
| best = None |
| for ra, rb in runs(sil[f_a:f_b]): |
| if best is None or rb - ra > best[1] - best[0]: |
| best = (ra, rb) |
| lo, hi = PAUSE_TARGET[words[i]["punct"]] |
| if best is None: |
| cur_sil = 0.0; mid = int((e_i + s_n) / 2 * sr) |
| else: |
| ra, rb = (f_a + best[0]) * hop, (f_a + best[1]) * hop |
| cur_sil = (rb - ra) / sr; mid = (ra + rb) // 2 |
| if cur_sil > hi and best is not None: |
| cut = int((cur_sil - hi) * sr) |
| c0 = mid - cut // 2 |
| pieces.append(y[cur:c0]); cur = c0 + cut |
| pauses.append(hi) |
| elif cur_sil < lo: |
| add = int((lo - cur_sil) * sr) |
| g0, g1 = int(e_i * sr), max(int(s_n * sr), int(e_i * sr) + hop) |
| seg = y[g0:g1] |
| if len(seg) >= hop: |
| k = int(np.argmin(frame_db(seg, sr))) * hop |
| ins_at = g0 + k |
| q = y[ins_at:ins_at + hop] |
| else: |
| ins_at, q = mid, np.zeros(hop, dtype=y.dtype) |
| filler = np.tile(q, add // hop + 1)[:add] * 0.5 |
| pieces.append(y[cur:ins_at]); pieces.append(filler); cur = ins_at |
| pauses.append(lo) |
| else: |
| pauses.append(cur_sil) |
| end = min(len(y), int(min(hi_s, ali[seg_b]["end"] + TAIL_S) * sr)) |
| pieces.append(y[cur:end]) |
| out = np.concatenate([p for p in pieces if len(p)]) if pieces else y[start:end] |
| |
| lead_have = a0 - start / sr |
| if lead_have < LEAD_S - 0.02: |
| out = np.concatenate([np.zeros(int((LEAD_S - lead_have) * sr), out.dtype), out]) |
| tail_have = end / sr - ali[seg_b]["end"] |
| if tail_have < TAIL_S - 0.02: |
| out = np.concatenate([out, np.zeros(int((TAIL_S - tail_have) * sr), out.dtype)]) |
| |
| f = min(xf, len(out) // 4) |
| if f > 0: |
| out[:f] *= np.linspace(0, 1, f); out[-f:] *= np.linspace(1, 0, f) |
| return out.astype(np.float32), pauses |
|
|
|
|
| |
| def pause_stats(y, sr, min_pause=0.15): |
| """VAD (energy, -40 dB vs 98th pct) pauses: leading/trailing silence and internal pauses >= min_pause.""" |
| sil = silence_mask(y, sr) |
| R = runs(sil) |
| n = len(sil) |
| lead = R[0][1] * 0.01 if R and R[0][0] == 0 else 0.0 |
| tail = (n - R[-1][0]) * 0.01 if R and R[-1][1] == n else 0.0 |
| internal = [(b - a) * 0.01 for a, b in R if a > 0 and b < n and (b - a) * 0.01 >= min_pause] |
| speech = max(1e-3, n * 0.01 - lead - tail) |
| return dict(lead=round(lead, 3), tail=round(tail, 3), pauses=[round(p, 3) for p in internal], |
| n_pauses=len(internal), pause_rate=round(len(internal) / speech, 4), |
| pause_frac=round(sum(internal) / speech, 4), speech_s=round(speech, 3)) |
|
|
|
|
| def f0_track(y16): |
| import pyworld as pw |
| x = y16.astype(np.float64) |
| f0, t = pw.dio(x, 16000, f0_floor=60, f0_ceil=500, frame_period=10) |
| return pw.stonemask(x, f0, t, 16000) |
|
|
|
|
| def f0_stats(y16): |
| f0 = f0_track(y16) |
| v = f0[f0 > 0] |
| if len(v) < 20: |
| return dict(voiced=float(len(v)), f0_med=None, st_std=None, st_range=None) |
| st = 12 * np.log2(v / np.median(v)) |
| d = np.diff(st) |
| return dict(voiced_frac=round(len(v) / len(f0), 3), f0_med=round(float(np.median(v)), 1), |
| st_std=round(float(st.std()), 3), st_range=round(float(np.percentile(st, 95) - np.percentile(st, 5)), 3), |
| st_vel=round(float(np.abs(d[np.abs(d) < 3]).mean()), 4)) |
|
|
|
|
| def tone_pairs(y16, words, ali): |
| """Yoruba tone-direction accuracy: for consecutive tone-bearing vowels whose tones differ, is the sign of the |
| median-F0 change (semitones) the same as the sign of the tone change? -> (n_pairs, n_correct)""" |
| f0 = f0_track(y16) |
| vals = [] |
| for w, a in zip(words, ali): |
| k = 0 |
| for (rc, oi, tone), (s, e) in zip([c for c in w["chars"] if c[0] in "abcdefghijklmnopqrstuvwxyz'"], a["chars"]): |
| if rc in _VOW or (rc in "nm" and tone is not None): |
| t = 1 if tone is None else tone |
| fr = f0[int(s / 0.01): int(e / 0.01) + 1] |
| fr = fr[fr > 0] |
| vals.append((t, float(np.median(12 * np.log2(fr))) if len(fr) else None)) |
| n = c = 0 |
| for (t1, f1), (t2, f2) in zip(vals, vals[1:]): |
| if t1 == t2 or f1 is None or f2 is None: |
| continue |
| n += 1 |
| c += int(np.sign(f2 - f1) == np.sign(t2 - t1)) |
| return n, c |
|
|
|
|
| def wasserstein1(a, b): |
| """1-D earth mover's distance between two samples.""" |
| if not len(a) or not len(b): |
| return None |
| a, b = np.sort(a), np.sort(b) |
| q = np.linspace(0, 1, 101) |
| return float(np.mean(np.abs(np.quantile(a, q) - np.quantile(b, q)))) |
|
|
