231 lines
10 KiB
Python
231 lines
10 KiB
Python
"""Forced aligner sidecar: where does each word of a known text start and end in its audio?
|
|
|
|
POST /align multipart: audio (wav/mp3/flac file), text, language, [start, end] seconds
|
|
POST /align/batch multipart: audio, spans (JSON [{start,end,text}]), language -- many sentences, one upload
|
|
GET /health
|
|
|
|
Times are seconds in the uploaded audio. Words are the whitespace-separated tokens of the text, in order; a
|
|
token with no letters (digits, punctuation) cannot be aligned and gets the gap between its neighbours
|
|
(aligned=false). Model: torchaudio MMS_FA (weights are fetched once into TORCH_HOME, a volume, never baked in).
|
|
No secrets, no auth: the service is reachable only on the internal Docker network.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import io
|
|
import json
|
|
import os
|
|
import re
|
|
import threading
|
|
import time
|
|
import unicodedata
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
|
|
|
|
SAMPLE_RATE = 16000
|
|
MAX_AUDIO_BYTES = int(os.environ.get("ALIGNER_MAX_AUDIO_BYTES", 40 * 1024 * 1024))
|
|
MAX_SPAN_SECONDS = float(os.environ.get("ALIGNER_MAX_SPAN_SECONDS", 120))
|
|
MAX_AUDIO_SECONDS = float(os.environ.get("ALIGNER_MAX_AUDIO_SECONDS", 900))
|
|
MAX_SPANS = int(os.environ.get("ALIGNER_MAX_SPANS", 400))
|
|
MAX_TEXT_CHARS = int(os.environ.get("ALIGNER_MAX_TEXT_CHARS", 20000))
|
|
REQUEST_TIMEOUT = float(os.environ.get("ALIGNER_REQUEST_TIMEOUT", 110))
|
|
QUEUE_WAIT = float(os.environ.get("ALIGNER_QUEUE_WAIT", 30))
|
|
DEVICE = os.environ.get("ALIGNER_DEVICE", "cpu")
|
|
_LANG_RE = re.compile(r"^[A-Za-z]{2,3}(-[A-Za-z0-9]{2,8})?$")
|
|
|
|
app = FastAPI(title="aligner", docs_url=None, redoc_url=None, openapi_url=None)
|
|
_gate = asyncio.Semaphore(1) # one alignment at a time: the model is the scarce resource
|
|
_state: dict[str, Any] = {"bundle": None, "model": None, "tok": None, "aligner": None, "error": None}
|
|
_load_lock = threading.Lock()
|
|
|
|
|
|
# ── text ───────────────────────────────────────────────────────────────────────
|
|
|
|
def normalize_word(word: str) -> str:
|
|
"""Romanize one token to the a-z' alphabet of the MMS_FA model (umlauts and accents folded)."""
|
|
w = unicodedata.normalize("NFKD", word.lower().replace("ß", "ss").replace("æ", "ae").replace("œ", "oe"))
|
|
w = "".join(c for c in w if not unicodedata.combining(c))
|
|
return re.sub(r"[^a-z']", "", w)
|
|
|
|
|
|
def tokenize(text: str) -> list[str]:
|
|
return text.split()
|
|
|
|
|
|
# ── model ──────────────────────────────────────────────────────────────────────
|
|
|
|
def _load() -> None:
|
|
with _load_lock:
|
|
if _state["model"] is not None:
|
|
return
|
|
import torch # imported lazily: /health works before the model is loaded
|
|
import torchaudio
|
|
torch.set_num_threads(int(os.environ.get("ALIGNER_THREADS", "4")))
|
|
bundle = torchaudio.pipelines.MMS_FA
|
|
model = bundle.get_model(with_star=False).to(DEVICE).eval()
|
|
_state.update(bundle=bundle, model=model, tok=bundle.get_tokenizer(), aligner=bundle.get_aligner(), error=None)
|
|
|
|
|
|
def align_samples(samples: np.ndarray, words: list[str]) -> list[tuple[float, float, bool]]:
|
|
"""Per whitespace token (start, end, aligned) in seconds relative to `samples` (16 kHz mono float32)."""
|
|
import torch
|
|
norm = [normalize_word(w) for w in words]
|
|
keep = [i for i, n in enumerate(norm) if n]
|
|
out: list[list[Any]] = [[None, None, False] for _ in words]
|
|
if keep and len(samples) >= 400:
|
|
wav = torch.from_numpy(samples)[None].to(DEVICE)
|
|
with torch.inference_mode():
|
|
emission, _ = _state["model"](wav)
|
|
spans = _state["aligner"](emission[0], _state["tok"]([norm[i] for i in keep]))
|
|
frame = len(samples) / emission.shape[1] / SAMPLE_RATE
|
|
for i, sp in zip(keep, spans):
|
|
out[i] = [sp[0].start * frame, sp[-1].end * frame, True]
|
|
# tokens without letters: the gap between their neighbours
|
|
total = len(samples) / SAMPLE_RATE
|
|
for i, o in enumerate(out):
|
|
if o[2]:
|
|
continue
|
|
prev_end = next((out[j][1] for j in range(i - 1, -1, -1) if out[j][2]), 0.0)
|
|
next_start = next((out[j][0] for j in range(i + 1, len(out)) if out[j][2]), total)
|
|
o[0], o[1] = prev_end, max(prev_end, next_start)
|
|
return [(float(a), float(b), bool(c)) for a, b, c in out]
|
|
|
|
|
|
# ── audio ──────────────────────────────────────────────────────────────────────
|
|
|
|
def decode(data: bytes) -> np.ndarray:
|
|
"""Any libsndfile format (wav, flac, mp3) to 16 kHz mono float32."""
|
|
import soundfile as sf
|
|
try:
|
|
with sf.SoundFile(io.BytesIO(data)) as f:
|
|
if f.frames / max(1, f.samplerate) > MAX_AUDIO_SECONDS:
|
|
raise HTTPException(413, f"audio longer than {MAX_AUDIO_SECONDS:.0f} s")
|
|
x = f.read(dtype="float32", always_2d=True)
|
|
sr = f.samplerate
|
|
except HTTPException:
|
|
raise
|
|
except Exception:
|
|
raise HTTPException(415, "audio is not a readable wav/mp3/flac file")
|
|
x = x.mean(axis=1)
|
|
if sr != SAMPLE_RATE:
|
|
import torch
|
|
import torchaudio.functional as F
|
|
x = F.resample(torch.from_numpy(x), sr, SAMPLE_RATE).numpy()
|
|
return np.ascontiguousarray(x, dtype=np.float32)
|
|
|
|
|
|
def _check_language(language: str) -> None:
|
|
if not _LANG_RE.match(language or ""):
|
|
raise HTTPException(422, "language must be an ISO code such as 'de' or 'en'")
|
|
|
|
|
|
def _check_text(text: str, strict: bool = True) -> list[str]:
|
|
if not isinstance(text, str) or len(text) > MAX_TEXT_CHARS:
|
|
raise HTTPException(422, f"text must be a string of at most {MAX_TEXT_CHARS} characters")
|
|
words = tokenize(text)
|
|
if not words:
|
|
raise HTTPException(422, "text is empty")
|
|
if strict and not any(normalize_word(w) for w in words):
|
|
raise HTTPException(422, "text has no letters the model can align (latin alphabets only)")
|
|
return words
|
|
|
|
|
|
def _slice(x: np.ndarray, start: float, end: float) -> tuple[np.ndarray, float]:
|
|
total = len(x) / SAMPLE_RATE
|
|
if not (0 <= start < end <= total + 0.05):
|
|
raise HTTPException(422, f"span {start}-{end} s is outside the audio (0-{total:.2f} s)")
|
|
if end - start > MAX_SPAN_SECONDS:
|
|
raise HTTPException(422, f"span longer than {MAX_SPAN_SECONDS:.0f} s")
|
|
return x[int(start * SAMPLE_RATE):int(min(end, total) * SAMPLE_RATE)], start
|
|
|
|
|
|
def _run(x: np.ndarray, spans: list[tuple[float, float, str]], deadline: float) -> list[list[dict[str, Any]]]:
|
|
_load()
|
|
result = []
|
|
for start, end, text in spans:
|
|
if time.monotonic() > deadline:
|
|
raise HTTPException(504, "alignment took longer than the request limit")
|
|
words = tokenize(text)
|
|
piece, offset = _slice(x, start, end)
|
|
times = align_samples(piece, words)
|
|
result.append([{"word": w, "start": round(offset + a, 3), "end": round(offset + b, 3), "aligned": ok}
|
|
for w, (a, b, ok) in zip(words, times)])
|
|
return result
|
|
|
|
|
|
async def _guarded(x: np.ndarray, spans: list[tuple[float, float, str]]) -> list[list[dict[str, Any]]]:
|
|
try:
|
|
await asyncio.wait_for(_gate.acquire(), QUEUE_WAIT)
|
|
except asyncio.TimeoutError:
|
|
raise HTTPException(503, "aligner busy, retry shortly", headers={"Retry-After": "10"})
|
|
try:
|
|
deadline = time.monotonic() + REQUEST_TIMEOUT
|
|
return await asyncio.wait_for(asyncio.to_thread(_run, x, spans, deadline), REQUEST_TIMEOUT + 5)
|
|
except asyncio.TimeoutError:
|
|
raise HTTPException(504, "alignment took longer than the request limit")
|
|
except HTTPException:
|
|
raise
|
|
except Exception as exc: # model download/load failure and the like
|
|
_state["error"] = type(exc).__name__
|
|
raise HTTPException(503, f"aligner model unavailable ({type(exc).__name__})")
|
|
finally:
|
|
_gate.release()
|
|
|
|
|
|
async def _read_audio(audio: UploadFile) -> np.ndarray:
|
|
data = await audio.read(MAX_AUDIO_BYTES + 1)
|
|
if len(data) > MAX_AUDIO_BYTES:
|
|
raise HTTPException(413, "audio file too large")
|
|
return await asyncio.to_thread(decode, data)
|
|
|
|
|
|
@app.on_event("startup")
|
|
async def _preload() -> None:
|
|
"""Load (and on first run download) the model in the background so the first request is not slow."""
|
|
def work() -> None:
|
|
try:
|
|
_load()
|
|
except Exception as exc: # reported by /health, retried on the first request
|
|
_state["error"] = type(exc).__name__
|
|
threading.Thread(target=work, name="preload", daemon=True).start()
|
|
|
|
|
|
# ── routes ─────────────────────────────────────────────────────────────────────
|
|
|
|
@app.get("/health")
|
|
async def health():
|
|
return {"status": "ok", "model_loaded": _state["model"] is not None, "device": DEVICE,
|
|
"last_error": _state["error"]}
|
|
|
|
|
|
@app.post("/align")
|
|
async def align(audio: UploadFile = File(...), text: str = Form(...), language: str = Form("de"),
|
|
start: float | None = Form(None), end: float | None = Form(None)):
|
|
_check_language(language)
|
|
_check_text(text)
|
|
x = await _read_audio(audio)
|
|
s = 0.0 if start is None else start
|
|
e = len(x) / SAMPLE_RATE if end is None else end
|
|
words = (await _guarded(x, [(s, e, text)]))[0]
|
|
return {"language": language, "words": words}
|
|
|
|
|
|
@app.post("/align/batch")
|
|
async def align_batch(audio: UploadFile = File(...), spans: str = Form(...), language: str = Form("de")):
|
|
_check_language(language)
|
|
try:
|
|
raw = json.loads(spans)
|
|
parsed = [(float(r["start"]), float(r["end"]), str(r["text"])) for r in raw]
|
|
except (ValueError, KeyError, TypeError):
|
|
raise HTTPException(422, "spans must be a JSON list of {start, end, text}")
|
|
if not parsed or len(parsed) > MAX_SPANS:
|
|
raise HTTPException(422, f"between 1 and {MAX_SPANS} spans are allowed")
|
|
for _, _, t in parsed:
|
|
_check_text(t, strict=False) # a sentence of digits only just gets its gap back
|
|
if sum(len(t) for _, _, t in parsed) > MAX_TEXT_CHARS * 4:
|
|
raise HTTPException(422, "too much text")
|
|
x = await _read_audio(audio)
|
|
return {"language": language, "spans": await _guarded(x, parsed)}
|