234 lines
10 KiB
Python
234 lines
10 KiB
Python
"""Word-level timing of a narrated article, measured by the forced-aligner sidecar (aligner/).
|
|
|
|
The MP3 is rendered sentence by sentence, so `audio/<id>.timings.json` already says where every sentence sits.
|
|
This module asks the sidecar where each *word* of a sentence starts and ends inside that span and records the
|
|
answer as `words: [[start_ms, end_ms], ...]` per segment (the file becomes version 2; version 1 readers simply
|
|
ignore the key). Never blocks playback: it runs in a background thread, retries with backoff, and when the
|
|
sidecar is down it gives up quietly -- the player then interpolates words inside each sentence as before.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import io
|
|
import json
|
|
import logging
|
|
import os
|
|
import threading
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Sequence
|
|
|
|
import requests
|
|
|
|
from core import epub_export
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
CHUNK_SECONDS = 45.0 # audio per sidecar request: bounded time per call and real progress
|
|
CHUNK_MAX_SPANS = 25
|
|
ATTEMPTS = 3
|
|
CONNECT_TIMEOUT = 3.0
|
|
READ_TIMEOUT = 150.0
|
|
FAILURE_COOLDOWN = 600.0 # after a failed run, do not retry on every timings request
|
|
WORDS_VERSION = 2
|
|
|
|
_jobs: dict[tuple[str, str], dict[str, Any]] = {}
|
|
_lock = threading.Lock()
|
|
|
|
|
|
def aligner_url() -> str:
|
|
return os.environ.get("ALIGNER_URL", "").strip().rstrip("/")
|
|
|
|
|
|
# ── planning ───────────────────────────────────────────────────────────────────
|
|
|
|
def alignment_jobs(units: Sequence[dict[str, Any]], blocks: Sequence[dict[str, Any]],
|
|
interview: bool) -> list[dict[str, Any]]:
|
|
"""Group sentence units into alignment jobs: [{units: [idx...], start, end, text, lead}] (ms).
|
|
|
|
A block whose sentences all have their own recorded utterance gets one job per sentence (tightest span).
|
|
A block whose spans were only anchored (counts disagreed) is aligned as one span and split afterwards.
|
|
`lead` counts leading tokens that were never spoken (an interview "Speaker:" label) and get no time.
|
|
"""
|
|
from core.magazine import _INTERVIEW_LABEL_RE
|
|
by_block: dict[int, list[int]] = {}
|
|
for idx, u in enumerate(units):
|
|
by_block.setdefault(u["block"], []).append(idx)
|
|
jobs: list[dict[str, Any]] = []
|
|
for block, idxs in by_block.items():
|
|
lead = 0
|
|
if interview:
|
|
m = _INTERVIEW_LABEL_RE.match(units[idxs[0]]["text"])
|
|
if m:
|
|
lead = len(units[idxs[0]]["text"].split()) - len(m.group(2).split())
|
|
exact = all(units[i].get("nth") is not None for i in idxs)
|
|
groups = [[i] for i in idxs] if exact else [idxs]
|
|
for g in groups:
|
|
text = " ".join(units[i]["text"] for i in g)
|
|
first_lead = lead if g[0] == idxs[0] else 0
|
|
spoken = " ".join(text.split()[first_lead:])
|
|
if not spoken:
|
|
continue
|
|
jobs.append({"units": g, "start": units[g[0]]["start"], "end": units[g[-1]]["end"],
|
|
"text": spoken, "lead": first_lead})
|
|
return jobs
|
|
|
|
|
|
def _chunks(jobs: Sequence[dict[str, Any]]) -> list[list[dict[str, Any]]]:
|
|
out: list[list[dict[str, Any]]] = []
|
|
cur: list[dict[str, Any]] = []
|
|
for j in jobs:
|
|
if cur and (len(cur) >= CHUNK_MAX_SPANS or (j["end"] - cur[0]["start"]) / 1000 > CHUNK_SECONDS):
|
|
out.append(cur)
|
|
cur = []
|
|
cur.append(j)
|
|
if cur:
|
|
out.append(cur)
|
|
return out
|
|
|
|
|
|
# ── sidecar client ─────────────────────────────────────────────────────────────
|
|
|
|
class AlignerError(RuntimeError):
|
|
"""The sidecar could not be used (down, busy too long, rejected the request)."""
|
|
|
|
|
|
def _post_chunk(url: str, wav: bytes, spans: list[dict[str, Any]], language: str,
|
|
sleep: Callable[[float], None] = time.sleep) -> list[list[dict[str, Any]]]:
|
|
last: Exception | None = None
|
|
for attempt in range(ATTEMPTS):
|
|
try:
|
|
r = requests.post(f"{url}/align/batch", files={"audio": ("chunk.wav", wav, "audio/wav")},
|
|
data={"spans": json.dumps(spans), "language": language},
|
|
timeout=(CONNECT_TIMEOUT, READ_TIMEOUT))
|
|
if r.status_code == 200:
|
|
return r.json()["spans"]
|
|
if 400 <= r.status_code < 500 and r.status_code != 429:
|
|
raise AlignerError(f"aligner rejected the request ({r.status_code})")
|
|
last = AlignerError(f"aligner answered {r.status_code}")
|
|
except requests.RequestException as exc:
|
|
last = AlignerError(f"aligner unreachable ({type(exc).__name__})")
|
|
except (ValueError, KeyError) as exc:
|
|
last = AlignerError(f"aligner sent an unreadable answer ({type(exc).__name__})")
|
|
if attempt < ATTEMPTS - 1:
|
|
sleep(2 ** attempt)
|
|
raise last or AlignerError("aligner failed")
|
|
|
|
|
|
def _wav_slice(audio: Any, start_ms: int, end_ms: int) -> bytes:
|
|
buf = io.BytesIO()
|
|
audio[start_ms:end_ms].set_frame_rate(16000).set_channels(1).export(buf, format="wav")
|
|
return buf.getvalue()
|
|
|
|
|
|
# ── the job ────────────────────────────────────────────────────────────────────
|
|
|
|
def align_words(mp3: Path, units: Sequence[dict[str, Any]], jobs: Sequence[dict[str, Any]], language: str, *,
|
|
url: str, progress: Callable[[int, int], None] | None = None,
|
|
post: Callable[..., Any] = _post_chunk) -> dict[int, list[list[int]]]:
|
|
"""Ask the sidecar for the words of every job; returns {unit index: [[start_ms, end_ms], ...]} for the
|
|
units it could time (a unit's list has one pair per whitespace token, label tokens first at zero length)."""
|
|
from pydub import AudioSegment
|
|
audio = AudioSegment.from_file(mp3)
|
|
chunks = _chunks(jobs)
|
|
result: dict[int, list[list[int]]] = {}
|
|
done = 0
|
|
for chunk in chunks:
|
|
lo = max(0, int(chunk[0]["start"]) - 100)
|
|
hi = min(len(audio), int(chunk[-1]["end"]) + 100)
|
|
spans = [{"start": round((j["start"] - lo) / 1000, 3), "end": round(min(j["end"], hi) / 1000 - lo / 1000, 3),
|
|
"text": j["text"]} for j in chunk]
|
|
answer = post(url, _wav_slice(audio, lo, hi), spans, language)
|
|
for j, words in zip(chunk, answer):
|
|
pairs = [[int(round(lo + w["start"] * 1000)), int(round(lo + w["end"] * 1000))] for w in words]
|
|
pairs = [[j["start"], j["start"]]] * j["lead"] + pairs
|
|
at = 0
|
|
for i in j["units"]:
|
|
n = len(units[i]["text"].split())
|
|
result[i] = pairs[at:at + n]
|
|
at += n
|
|
done += 1
|
|
if progress:
|
|
progress(done, len(jobs))
|
|
return result
|
|
|
|
|
|
def build_segment_words(units: Sequence[dict[str, Any]], per_unit: dict[int, list[list[int]]],
|
|
n_segments: int) -> dict[int, list[list[int]]]:
|
|
"""Per record segment, the words of all its sentences in order -- only for segments fully covered."""
|
|
out: dict[int, list[list[int]]] = {}
|
|
bad: set[int] = set()
|
|
for idx, u in enumerate(units):
|
|
k = u["seg"]
|
|
words = per_unit.get(idx)
|
|
if words is None or len(words) != len(u["text"].split()):
|
|
bad.add(k)
|
|
continue
|
|
out.setdefault(k, []).extend(words)
|
|
return {k: v for k, v in out.items() if k not in bad and 0 <= k < n_segments}
|
|
|
|
|
|
def write_words(mp3: Path, segment_words: dict[int, list[list[int]]]) -> bool:
|
|
"""Add `words` to the record (atomic replace; everything else in the file is kept as it was)."""
|
|
path = mp3.with_suffix(".timings.json")
|
|
try:
|
|
rec = json.loads(path.read_text("utf-8"))
|
|
segs = rec["segments"]
|
|
for k, words in segment_words.items():
|
|
segs[k]["words"] = words
|
|
rec["version"] = WORDS_VERSION
|
|
rec["words_engine"] = "mms_fa"
|
|
tmp = path.with_suffix(".tmp")
|
|
tmp.write_text(json.dumps(rec), "utf-8")
|
|
tmp.replace(path)
|
|
return True
|
|
except (OSError, ValueError, KeyError, IndexError):
|
|
log.warning("Could not write word timings for %s", mp3.name)
|
|
return False
|
|
|
|
|
|
# ── background job registry ────────────────────────────────────────────────────
|
|
|
|
def status(doc_id: str, article_id: str) -> dict[str, Any]:
|
|
with _lock:
|
|
return dict(_jobs.get((doc_id, article_id)) or {"state": "idle", "done": 0, "total": 0})
|
|
|
|
|
|
def should_autostart(doc_id: str, article_id: str) -> bool:
|
|
if not aligner_url():
|
|
return False
|
|
st = status(doc_id, article_id)
|
|
if st["state"] == "running":
|
|
return False
|
|
# Whatever the last outcome, do not re-run on every request: a run that finished without covering every
|
|
# sentence (or after the audio was re-rendered) is retried after the cooldown.
|
|
return time.time() - st.get("finished", 0) >= FAILURE_COOLDOWN
|
|
|
|
|
|
def start(doc_id: str, article_id: str, run: Callable[[Callable[[int, int], None]], dict[str, Any]]) -> bool:
|
|
"""Run `run(progress)` in a daemon thread unless this article is already being aligned."""
|
|
key = (doc_id, article_id)
|
|
with _lock:
|
|
if (_jobs.get(key) or {}).get("state") == "running":
|
|
return False
|
|
_jobs[key] = {"state": "running", "done": 0, "total": 0, "started": time.time()}
|
|
|
|
def progress(done: int, total: int) -> None:
|
|
with _lock:
|
|
_jobs[key].update(done=done, total=total)
|
|
|
|
def target() -> None:
|
|
try:
|
|
info = run(progress)
|
|
end = {"state": "done", **info}
|
|
except AlignerError as exc:
|
|
log.warning("Word alignment for %s skipped: %s", article_id, exc)
|
|
end = {"state": "failed", "error": str(exc)}
|
|
except Exception as exc: # never let a worker thread die noisily
|
|
log.exception("Word alignment for %s failed", article_id)
|
|
end = {"state": "failed", "error": type(exc).__name__}
|
|
with _lock:
|
|
_jobs[key].update(end, finished=time.time())
|
|
|
|
threading.Thread(target=target, name=f"align-{article_id}", daemon=True).start()
|
|
return True
|