ApiCore/ApiCore-Whisper/transcribe.py

"""Batch transcription for the voice CronJob (Task 13, spec 6.3 and 9.2).

Second init container of ``apicore-job-voice-transcribe``. It reads
``<work>/in/manifest.json`` written by the ``fetch`` phase, transcribes the
listed audio files one after the other with one loaded model and writes one
``<work>/out/<captureId>.json`` per file for the ``ingest`` phase.

The container has no secrets, no database and no network. Belt and braces to
the NetworkPolicy: ``socket.socket`` is blocked before the model loads.
Transcripts and audio are never printed or logged; stdout carries counters
only, stderr only exception class names.

Result file (matches ``WhisperResult`` in the backend ``work-dir.ts``)::

    {v: 1, captureId, text, durationSec, language, languageProbability,
     segments: [{start, end, text}], error?: NO_SPEECH | TOO_LONG | DECODE_FAILED}
"""

from __future__ import annotations

import json
import os
import re
import socket
import stat
import sys
from collections.abc import Callable, Iterable
from pathlib import Path
from typing import Any

RESULT_VERSION = 1
MAX_AUDIO_SEC = 240.0
MIN_TEXT_CHARS = 5
GERMAN = "de"
MIN_FOREIGN_LANGUAGE_PROBABILITY = 0.6
MAX_AUDIO_BYTES = 3 * 1024 * 1024
MAX_MANIFEST_BYTES = 256 * 1024
MAX_ITEMS = 50
DEFAULT_WORK_DIR = "/work"
DEFAULT_MODEL_DIR = "/models/small"
DEFAULT_CPU_THREADS = 2
MAX_CPU_THREADS = 16
UNKNOWN_LANGUAGE = "und"
SAMPLE_RATE = 16000
# Container per file extension (the manifest only allows these three). Naming
# the demuxer keeps libav from probing every demuxer on untrusted bytes.
CONTAINER_FORMATS = {"webm": "matroska", "ogg": "ogg", "mp4": "mov"}

ID_PATTERN = re.compile(r"^[A-Za-z0-9_-]{1,64}$")
FILE_PATTERN = re.compile(r"^([A-Za-z0-9_-]{1,64})\.(webm|ogg|mp4)$")


class DecodeFailed(Exception):
    """The file is not usable audio. Carries no content."""


class ManifestInvalid(Exception):
    """The manifest exists but is not valid."""


class _BlockedSocket(socket.socket):
    """A socket that cannot be created (a class, so later subclassing works)."""

    def __init__(self, *args: Any, **kwargs: Any) -> None:  # noqa: D107
        raise OSError("network access is disabled in the whisper container")


def block_network() -> None:
    """Make every later ``socket.socket(...)`` fail."""
    socket.socket = _BlockedSocket  # type: ignore[misc,assignment]


# --- audio ---------------------------------------------------------------


def container_format(path: Path) -> str:
    try:
        return CONTAINER_FORMATS[path.suffix.lstrip(".")]
    except KeyError:
        raise DecodeFailed("unsupported container") from None


def decode_audio(path: Path, limit: float = MAX_AUDIO_SEC) -> Any:
    """Mono float32 samples at 16 kHz (what Whisper takes), decoded here so the
    demuxer is named explicitly. Stops a little past ``limit`` seconds."""
    import av
    import numpy as np

    max_samples = int((limit + 1) * SAMPLE_RATE)
    try:
        with av.open(str(path), mode="r", format=container_format(path)) as container:
            stream = next((s for s in container.streams if s.type == "audio"), None)
            if stream is None:
                raise DecodeFailed("no audio stream")
            resampler = av.AudioResampler(format="s16", layout="mono", rate=SAMPLE_RATE)
            chunks = []
            total = 0
            for frame in container.decode(stream):
                for out in resampler.resample(frame):
                    chunk = out.to_ndarray().reshape(-1)
                    chunks.append(chunk)
                    total += chunk.size
                if total > max_samples:
                    break
            for out in resampler.resample(None):
                chunks.append(out.to_ndarray().reshape(-1))
    except DecodeFailed:
        raise
    except Exception as error:
        raise DecodeFailed(type(error).__name__) from None
    if not chunks:
        raise DecodeFailed("no samples")
    return np.concatenate(chunks).astype(np.float32) / 32768.0


def measure_duration(path: Path, limit: float = MAX_AUDIO_SEC) -> float:
    """Measured length in seconds, by demuxing packets (no decoding).

    Stops as soon as the length passes ``limit`` (the value is then a lower
    bound above the limit). The container header is not trusted: browser
    recordings often have none, and a forged one must not hide a long file.
    """
    import av  # imported here so an idle run (no manifest) stays fast

    try:
        with av.open(str(path), mode="r", format=container_format(path)) as container:
            stream = next((s for s in container.streams if s.type == "audio"), None)
            if stream is None:
                raise DecodeFailed("no audio stream")
            end = 0.0
            counted = 0.0
            for packet in container.demux(stream):
                if packet.size == 0 or packet.pts is None or packet.time_base is None:
                    if packet.duration is not None and packet.time_base is not None:
                        counted += float(packet.duration * packet.time_base)
                        end = max(end, counted)
                    continue
                packet_end = float((packet.pts + (packet.duration or 0)) * packet.time_base)
                end = max(end, packet_end)
                counted = end
                if end > limit:
                    break
            if end <= 0.0:
                raise DecodeFailed("no measurable duration")
            return round(end, 2)
    except DecodeFailed:
        raise
    except Exception as error:  # av.error.FFmpegError, OSError, ValueError
        raise DecodeFailed(type(error).__name__) from None


# --- results -------------------------------------------------------------


def error_result(capture_id: str, error: str, duration: float = 0.0) -> dict[str, Any]:
    return {
        "v": RESULT_VERSION,
        "captureId": capture_id,
        "text": "",
        "durationSec": duration,
        "language": UNKNOWN_LANGUAGE,
        "languageProbability": 0.0,
        "segments": [],
        "error": error,
    }


def _clamp_probability(value: Any) -> float:
    try:
        number = float(value)
    except (TypeError, ValueError):
        return 0.0
    if number != number:  # NaN
        return 0.0
    return round(min(max(number, 0.0), 1.0), 3)


def build_result(
    capture_id: str, duration: float, segments: Iterable[Any], info: Any
) -> dict[str, Any]:
    parts = []
    for segment in segments:
        text = str(segment.text).strip()
        if text:
            parts.append(
                {
                    "start": round(float(segment.start), 2),
                    "end": round(float(segment.end), 2),
                    "text": text,
                }
            )
    full_text = " ".join(part["text"] for part in parts)
    language = str(getattr(info, "language", "") or "")
    if not re.fullmatch(r"[A-Za-z-]{1,16}", language):
        language = UNKNOWN_LANGUAGE
    probability = _clamp_probability(getattr(info, "language_probability", 0.0))
    foreign = language != GERMAN and probability < MIN_FOREIGN_LANGUAGE_PROBABILITY
    if len(full_text) < MIN_TEXT_CHARS or foreign:
        result = error_result(capture_id, "NO_SPEECH", duration)
        result["language"] = language
        result["languageProbability"] = probability
        return result
    return {
        "v": RESULT_VERSION,
        "captureId": capture_id,
        "text": full_text,
        "durationSec": duration,
        "language": language,
        "languageProbability": probability,
        "segments": parts,
    }


def write_result(out_dir: Path, capture_id: str, result: dict[str, Any]) -> None:
    """Writes by rename, mode 0600, so ingest never sees a half-written file."""
    target = out_dir / f"{capture_id}.json"
    temporary = out_dir / f"{capture_id}.json.tmp"
    flags = os.O_WRONLY | os.O_CREAT | os.O_TRUNC
    descriptor = os.open(temporary, flags, 0o600)
    with os.fdopen(descriptor, "w", encoding="utf-8") as handle:
        json.dump(result, handle, ensure_ascii=False)
    os.replace(temporary, target)


# --- manifest --------------------------------------------------------------


def load_manifest(work_dir: Path) -> list[dict[str, str]] | None:
    """Items of this run; None if there is no manifest (nothing to do)."""
    path = work_dir / "in" / "manifest.json"
    try:
        info = path.lstat()
    except FileNotFoundError:
        return None
    if not stat.S_ISREG(info.st_mode) or info.st_size > MAX_MANIFEST_BYTES:
        raise ManifestInvalid("manifest is not a regular file or too large")
    try:
        parsed = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, ValueError):
        raise ManifestInvalid("manifest is not readable JSON") from None
    items = parsed.get("items") if isinstance(parsed, dict) else None
    if not isinstance(items, list) or parsed.get("v") != 1 or len(items) > MAX_ITEMS:
        raise ManifestInvalid("manifest has an unexpected shape")
    checked: list[dict[str, str]] = []
    for item in items:
        capture_id = item.get("captureId") if isinstance(item, dict) else None
        file_name = item.get("file") if isinstance(item, dict) else None
        match = FILE_PATTERN.match(file_name) if isinstance(file_name, str) else None
        if not isinstance(capture_id, str) or not ID_PATTERN.match(capture_id):
            raise ManifestInvalid("manifest item without a usable id")
        if match is None or match.group(1) != capture_id:
            raise ManifestInvalid("manifest item with a file name that is not its id")
        checked.append({"captureId": capture_id, "file": file_name})
    return checked


# --- run -------------------------------------------------------------------

Counters = dict[str, int]
ModelFactory = Callable[[], Any]


def _usable_input(path: Path) -> bool:
    try:
        info = path.lstat()
    except OSError:
        return False
    return stat.S_ISREG(info.st_mode) and 0 < info.st_size <= MAX_AUDIO_BYTES


def transcribe_one(model: Any, path: Path, capture_id: str, prompt: str) -> dict[str, Any]:
    """One file to one result. Never raises: a broken file is DECODE_FAILED."""
    if not _usable_input(path):
        return error_result(capture_id, "DECODE_FAILED")
    try:
        duration = measure_duration(path)
    except DecodeFailed:
        return error_result(capture_id, "DECODE_FAILED")
    if duration > MAX_AUDIO_SEC:
        return error_result(capture_id, "TOO_LONG", duration)
    try:
        audio = decode_audio(path)
        if audio.size / SAMPLE_RATE > MAX_AUDIO_SEC:  # decoded length beats any header
            return error_result(capture_id, "TOO_LONG", round(audio.size / SAMPLE_RATE, 2))
        segments, info = model.transcribe(
            audio,
            language=GERMAN,
            beam_size=1,
            vad_filter=True,
            condition_on_previous_text=False,
            initial_prompt=prompt,
        )
        collected = list(segments)  # the generator does the work here
        return build_result(capture_id, duration, collected, info)
    except DecodeFailed:
        return error_result(capture_id, "DECODE_FAILED", duration)
    except Exception as error:  # decoder or model error on this file only
        print(f"transcribe failed: {type(error).__name__}", file=sys.stderr)
        return error_result(capture_id, "DECODE_FAILED", duration)


def run(
    work_dir: Path,
    model_factory: ModelFactory,
    prompt: str,
    items: list[dict[str, str]],
) -> Counters:
    counters: Counters = {"transcribed": 0, "no_speech": 0, "too_long": 0, "failed": 0}
    out_dir = work_dir / "out"
    out_dir.mkdir(mode=0o700, exist_ok=True)
    model = model_factory()
    for item in items:  # one after the other: the pod has one CPU budget
        result = transcribe_one(model, work_dir / "in" / item["file"], item["captureId"], prompt)
        write_result(out_dir, item["captureId"], result)
        error = result.get("error")
        key = {
            None: "transcribed",
            "NO_SPEECH": "no_speech",
            "TOO_LONG": "too_long",
        }.get(error, "failed")
        counters[key] += 1
    return counters


def cpu_threads(env: dict[str, str]) -> int:
    raw = env.get("WHISPER_CPU_THREADS", "")
    if not raw:
        return DEFAULT_CPU_THREADS
    if not raw.isdigit() or not 1 <= int(raw) <= MAX_CPU_THREADS:
        raise ValueError(f"WHISPER_CPU_THREADS must be a whole number from 1 to {MAX_CPU_THREADS}")
    return int(raw)


def load_whisper_model(model_dir: str, threads: int) -> Any:
    """Imports faster-whisper, then cuts the network, then loads the model."""
    from faster_whisper import WhisperModel

    block_network()
    return WhisperModel(model_dir, device="cpu", compute_type="int8", cpu_threads=threads)


def main(env: dict[str, str] | None = None) -> int:
    env = dict(os.environ) if env is None else env
    work_dir = Path(env.get("VOICE_WORK_DIR") or DEFAULT_WORK_DIR)
    try:
        items = load_manifest(work_dir)
    except ManifestInvalid as error:
        print(f"whisper: {error}", file=sys.stderr)
        return 1
    if not items:
        print("transcribed=0 no_speech=0 too_long=0 failed=0")
        return 0
    prompt_file = Path(env.get("WHISPER_VOCABULARY") or Path(__file__).with_name("vocabulary.txt"))
    prompt = prompt_file.read_text(encoding="utf-8").strip()
    model_dir = env.get("WHISPER_MODEL_DIR") or DEFAULT_MODEL_DIR
    threads = cpu_threads(env)
    counters = run(work_dir, lambda: load_whisper_model(model_dir, threads), prompt, items)
    print(" ".join(f"{name}={count}" for name, count in counters.items()))
    return 0


if __name__ == "__main__":
    sys.exit(main())