mathly / app /align.py
3v324v23's picture
Add keystore for managing API keys securely and update tests
35394c4
Raw History Blame Contribute Delete
15.5 kB
"""CTC forced alignment and goodness-of-pronunciation (GOP), in plain numpy.
Inputs are frame-level log-posteriors `logp` of shape (T, V) from a CTC phoneme model.
Nothing here depends on torch, so it is unit-testable with synthetic emissions.
"""
from __future__ import annotations
import logging
import os
from dataclasses import dataclass, field
import numpy as np
log = logging.getLogger("ihki.align")
NEG = -1e30
POOLS = ("max", "mean", "topk", "lse", "energy", "gated")
DEFAULT_POOL = "gated"
def env_pool() -> str:
"""The ratio pool GOP_RATIO_POOL names (read when a GopConfig is made, so after main.py loads .env).
Empty or unset: DEFAULT_POOL. An unknown name is logged and DEFAULT_POOL is used."""
pool = os.environ.get("GOP_RATIO_POOL", "").strip().lower() or DEFAULT_POOL
if pool not in POOLS:
log.warning("GOP_RATIO_POOL=%r is not one of %s: using %s", pool, ", ".join(POOLS), DEFAULT_POOL)
return DEFAULT_POOL
return pool
NARROW_PAD = 1
def env_vowel_window() -> int | None:
"""GOP_VOWEL_WINDOW: empty, unset or "wide" (None: a vowel reads the consonants' wide window), "narrow" (its
aligned span +- NARROW_PAD frames), or that pad as a frame count. Anything else is logged and "wide" used.
Wide is the default: with a pooled ratio it absorbs a transition spike without the narrow window's blind spot
(docs/gop-pooling.md, scripts/gop_window_bench.py)."""
raw = os.environ.get("GOP_VOWEL_WINDOW", "").strip().lower()
if raw in ("", "wide"):
return None
if raw == "narrow":
return NARROW_PAD
if raw.isdigit():
return int(raw)
log.warning("GOP_VOWEL_WINDOW=%r is not wide, narrow or a frame count: using wide", raw)
return None
@dataclass
class Alignment:
total_logp: float # log-prob of the best CTC path
spans: list[tuple[int, int]] # [start, end) frames where each target token is emitted
def ctc_viterbi(logp: np.ndarray, tokens: list[int], blank: int) -> Alignment | None:
"""Best CTC path that emits exactly `tokens`. Returns None if the audio is too short."""
T, _ = logp.shape
N = len(tokens)
if N == 0 or T < N:
return None
S = 2 * N + 1
ext = np.array([blank if s % 2 == 0 else tokens[s // 2] for s in range(S)])
# skip transition s-2 -> s allowed for non-blank states whose label differs from s-2
can_skip = np.zeros(S, dtype=bool)
can_skip[2:] = (ext[2:] != blank) & (ext[2:] != ext[:-2])
dp = np.full((T, S), NEG)
bp = np.zeros((T, S), dtype=np.int8) # 0 = stay, 1 = from s-1, 2 = from s-2
dp[0, 0] = logp[0, blank]
dp[0, 1] = logp[0, ext[1]]
for t in range(1, T):
prev = dp[t - 1]
stay = prev
one = np.concatenate(([NEG], prev[:-1]))
two = np.where(can_skip, np.concatenate(([NEG, NEG], prev[:-2])), NEG)
stacked = np.stack([stay, one, two])
choice = stacked.argmax(0)
dp[t] = stacked[choice, np.arange(S)] + logp[t, ext]
bp[t] = choice
# path must end in the last token or the trailing blank
end = S - 1 if dp[T - 1, S - 1] >= dp[T - 1, S - 2] else S - 2
total = float(dp[T - 1, end])
if total <= NEG / 2:
return None
states = np.empty(T, dtype=np.int64)
s = end
for t in range(T - 1, -1, -1):
states[t] = s
s -= int(bp[t, s])
spans: list[list[int]] = [[-1, -1] for _ in range(N)]
for t, st in enumerate(states):
if st % 2 == 1:
i = st // 2
if spans[i][0] < 0:
spans[i][0] = t
spans[i][1] = t + 1
return Alignment(total, [tuple(x) for x in spans])
@dataclass
class TokenScore:
score: float # 0..100
gop: float # log ratio target vs strongest competitor, <= 0
peak: float # max log-posterior of the target inside its window
competitor: int | None # strongest non-target, non-blank token in the window (None: nothing, or a sound we do not name)
deleted: bool # target barely present and nothing else either
realised: int | None = None # a class target (several correct realisations): the one said
other: bool = False # the strongest competitor is a sound outside the named set
comp_peak: float = NEG # max log-posterior of the strongest competitor (named or not) in the window
comp_frames: int = 0 # frames where that competitor clears deleted_comp (duration evidence)
window: tuple[int, int] = (0, 0) # [start, end) frames the target and its competitors were read from
@dataclass
class GopConfig:
deleted_peak: float = float(np.log(0.05)) # target this weak ...
deleted_comp: float = float(np.log(0.20)) # ... and no strong competitor -> deletion
confident_peak: float = 0.30 # target posterior at which the peak no longer lowers the score
# Softens the target/competitor ratio: at 0.5 a competitor twice as likely still scores 71, not 50.
# The zero-shot model is noisy on accented speech, so a linear ratio punishes its own uncertainty.
leniency: float = 0.5
# How the target/competitor ratio pools the window's frames (docs/gop-pooling.md). The peak behind
# `confidence` and `deleted` stays the max, so a peaky CTC emission is never diluted by its blanks.
ratio_pool: str = field(default_factory=env_pool) # max | mean | topk | lse | energy | gated (GOP_RATIO_POOL)
pool_k: int = 2 # topk: frames averaged per column
pool_tau: float = 0.5 # lse: temperature (0 -> max, 1 -> mean)
energy_gamma: float = 1.0 # energy: frame weight = (RMS / loudest in the window) ** gamma
energy_gate: float = 0.3 # gated: frames quieter than this fraction of the window's loudest are dropped
energy_spread: int = 1 # frames the energy is max-filtered over (CTC peaks lead/lag the acoustics)
pool_vowels_only: bool = True # consonants keep the max (bursts and voiceless sounds are short and quiet)
# Consonants need some surrounding blank for a CTC peak that leads/lags the forced span, but must not inherit an
# arbitrarily long pause up to the next token (e.g. a 760 ms window across a sentence break).
consonant_window: int | None = 4
# A vowel's competitors are read from its aligned span +- this many frames (None: the wide window consonants
# use). Its presence (`peak`, `deleted`) is still read from the wide window. GOP_VOWEL_WINDOW sets it.
vowel_window: int | None = field(default_factory=env_vowel_window)
def pool_frames(seg: np.ndarray, how: str = "max", k: int = 2, tau: float = 0.5,
weight: np.ndarray | None = None) -> np.ndarray:
"""Pool a window's (n, V) log-posteriors over time into one log-score per column.
max: the single best frame. topk: log of the mean of each column's k best frames. lse: the power mean
tau * log(mean(p ** (1/tau))), between the max (tau -> 0) and the mean (tau = 1). mean / energy / gated:
log of the weighted mean posterior, `weight` (uniform when None) normalised over the window.
"""
n = len(seg)
if how == "max" or n == 1:
return seg.max(0)
if how == "topk":
k = max(1, min(k, n))
return np.logaddexp.reduce(np.partition(seg, n - k, axis=0)[n - k:], axis=0) - np.log(k)
if how == "lse":
return tau * (np.logaddexp.reduce(seg / tau, axis=0) - np.log(n))
w = np.ones(n) if weight is None else np.asarray(weight, dtype=np.float64)
w = w / w.sum() if w.sum() > 0 else np.full(n, 1.0 / n)
return np.logaddexp.reduce(seg + np.log(np.maximum(w, 1e-300))[:, None], axis=0)
def energy_weights(energy: np.ndarray, how: str, gamma: float = 1.0, gate: float = 0.3) -> np.ndarray:
"""Frame weights of one window from its (already smoothed) RMS: soft (`energy`) or a 0/1 voice gate (`gated`)."""
e = energy / max(float(energy.max()), 1e-12)
return (e >= gate).astype(np.float64) if how == "gated" else e ** gamma
def spread_energy(energy: np.ndarray, spread: int) -> np.ndarray:
"""Max-filter over +-spread frames: a CTC emission a frame off the loud nucleus still gets its weight."""
out = energy.astype(np.float64).copy()
for s in range(1, spread + 1):
out[s:] = np.maximum(out[s:], energy[:-s])
out[:-s] = np.maximum(out[:-s], energy[s:])
return out
def token_window(spans: list[tuple[int, int]], i: int, T: int, pad: int | None = None) -> tuple[int, int]:
"""[start, end) frames token i is read from. Wide (pad None): from the end of token i-1 to the start of token i+1,
the blanks around the emission where a substituted consonant peaks. Narrow: its own span +- `pad` frames, inside
the wide window (a vowel is its syllable's loud nucleus, not a sound hiding in the blanks)."""
a = spans[i - 1][1] if i > 0 else 0
b = spans[i + 1][0] if i + 1 < len(spans) else T
a, b = min(a, spans[i][0]), max(b, spans[i][1])
if pad is None:
return a, b
return max(a, spans[i][0] - pad), min(b, spans[i][1] + pad)
def compute_token_gop(logp: np.ndarray, spans: list[tuple[int, int]], i: int, tokens: list[int], blank: int,
cfg: GopConfig, vowel: bool = False, accept: list[set[int] | None] | None = None,
n_real: int | None = None, named: set[int] | None = None, energy: np.ndarray | None = None,
edge: tuple[set[int], set[int]] | None = None, pooled: bool | None = None) -> TokenScore:
"""GOP of token i (see score_tokens), its window and pooling chosen by its class.
Presence is read from the wide window for every token: `peak` (behind the confidence) and `deleted` (target under
deleted_peak and no competitor over deleted_comp), so a vowel aligned to a single frame is still found deleted.
The competitor ratio of a vowel is read from its narrow window (cfg.vowel_window) and pooled with cfg.ratio_pool;
a consonant is capped at its span +- cfg.consonant_window and keeps the max unless cfg.pool_vowels_only is off.
`pooled` overrides the latter.
`edge`: ids that are no competitor on the frames between the token's span and its left / right neighbour (a vowel
gliding into j or w passes through i or u: coarticulation, not a substitution). `energy` is already spread.
"""
T = logp.shape[0]
accept = accept or [None] * len(tokens)
a, b = token_window(spans, i, T, cfg.consonant_window) if not vowel else token_window(spans, i, T)
ra, rb = token_window(spans, i, T, cfg.vowel_window) if vowel and cfg.vowel_window is not None else (a, b)
tok, (s0, s1) = tokens[i], spans[i]
seg = logp[a:b]
peak = float(seg[:, tok].max())
members = accept[i] or set()
realised = max(members, key=lambda m: seg[:, m].max()) if members else None
cols = (seg[:, :n_real] if n_real else seg).copy()
if edge:
for ids, lo, hi in ((edge[0], a, s0), (edge[1], s1, b)):
ids = [x for x in ids if x < cols.shape[1]]
if ids and hi > lo:
cols[lo - a:hi - a, ids] = NEG
mask = np.zeros(cols.shape[1], dtype=bool)
mask[blank] = True
for j in (i - 1, i, i + 1):
if 0 <= j < len(tokens):
for x in {tokens[j], *(accept[j] or ())}:
if x < len(mask):
mask[x] = True
deleted = bool(peak < cfg.deleted_peak and float(np.where(mask, NEG, cols.max(0)).max()) < cfg.deleted_comp)
how = cfg.ratio_pool if (pooled if pooled is not None else vowel or not cfg.pool_vowels_only) else "max"
weight = None
if how in ("energy", "gated"):
if energy is not None:
weight = energy_weights(energy[ra:rb], how, cfg.energy_gamma, cfg.energy_gate)
weight[max(0, s0 - ra):max(0, s1 - ra)] = 1.0 # the frames the token was emitted on count in full
how = "mean"
rcols = cols[ra - a:rb - a]
maxed = rcols.max(0)
ratio = maxed.copy() if how == "max" else pool_frames(rcols, how, cfg.pool_k, cfg.pool_tau, weight)
ratio[mask] = NEG
comp = int(ratio.argmax())
comp_peak = float(maxed[comp]) if ratio[comp] > NEG / 2 else NEG
# A lone CTC spike can be a breath or microphone onset rather than a produced consonant. Keep duration
# evidence separately from the peak so pronunciation policy can distinguish transients from substitutions.
comp_frames = int(np.count_nonzero(rcols[:, comp] >= cfg.deleted_comp)) if ratio[comp] > NEG / 2 else 0
tseg = logp[ra:rb, [tok]]
target = float(tseg.max()) if how == "max" else float(pool_frames(tseg, how, cfg.pool_k, cfg.pool_tau, weight)[0])
gop = target - max(target, float(ratio[comp]))
confidence = min(1.0, float(np.exp(peak)) / cfg.confident_peak)
score = 100.0 * float(np.exp(gop * cfg.leniency)) * float(np.sqrt(confidence))
present = comp_peak > NEG / 2
other = bool(present and named is not None and comp not in named)
return TokenScore(round(score, 1), round(gop, 3), round(peak, 3), comp if present and not other else None,
deleted, realised, other, round(comp_peak, 3), comp_frames, (int(ra), int(rb)))
def score_tokens(logp: np.ndarray, align: Alignment, tokens: list[int], blank: int,
cfg: GopConfig | None = None, accept: list[set[int] | None] | None = None,
n_real: int | None = None, named: set[int] | None = None,
energy: np.ndarray | None = None, pooled: list[bool] | None = None,
vowels: list[bool] | None = None,
edges: list[tuple[set[int], set[int]] | None] | None = None) -> list[TokenScore]:
"""Per-token GOP (compute_token_gop for each token).
The neighbours' own labels are excluded from the competitor set (coarticulation is not an error).
A class token (a virtual column pooling its correct realisations, `accept[i]` their ids) is scored on the
pooled column; its members are not competitors. Competitors come from the first `n_real` columns only, and
one outside `named` is reported as `other` rather than by name.
`vowels` marks the tokens that are vowels (narrow window, pooled ratio). Without it every token is read from the
wide window and pooled, unless `pooled` marks which are. `energy` is the per-frame RMS the energy pools need;
without it they fall back to the plain mean. `edges[i]`: see compute_token_gop.
"""
cfg = cfg or GopConfig()
if energy is not None:
energy = spread_energy(np.asarray(energy)[:logp.shape[0]], cfg.energy_spread)
return [compute_token_gop(logp, align.spans, i, tokens, blank, cfg, bool(vowels and vowels[i]), accept, n_real,
named, energy, edges[i] if edges else None,
(pooled[i] if pooled is not None else None if vowels is not None else True))
for i in range(len(tokens))]
def greedy_decode(logp: np.ndarray, blank: int) -> list[int]:
ids = logp.argmax(1)
out, prev = [], None
for x in ids:
x = int(x)
if x != prev and x != blank:
out.append(x)
prev = x
return out
def edit_distance(a: list, b: list, same=None) -> int:
same = same or (lambda x, y: x == y)
d = list(range(len(b) + 1))
for i, x in enumerate(a, 1):
prev, d[0] = d[0], i
for j, y in enumerate(b, 1):
prev, d[j] = d[j], min(d[j] + 1, d[j - 1] + 1, prev + (not same(x, y)))
return d[len(b)]