Blame for src/nida/data/workers/__init__.py

"""Phase 1 model workers (docs/03-data-pipeline.md, research note R8).

Each module here is one model family run in its own virtual environment
(:mod:`nida.data.envs`) as ``python -m nida.data.workers.<name> ...``. They reuse the
Phase 0 JSONL protocol and generic loop (:mod:`nida.bench.protocol`,
:mod:`nida.bench.workers.runner`): one ``WorkRequest`` in, one ``WorkResult`` out, heavy
imports only inside ``Adapter.load``.

Array outputs (frame-level probabilities, CTC emissions, ...) never travel inside the JSON
result. Every request names its output file in ``params["out_path"]`` (a ``.npy`` path); the
worker writes a float32 array there atomically and returns the path plus the facts a
consumer needs to interpret it (``num_frames``, ``frame_hop_ms``, ...) in the payload.

This ``__init__`` holds the small helpers shared by all four workers: request/audio
validation, the frame-count check that guards against a wrong hop assumption, and the atomic
``.npy`` writer. It imports only the standard library (numpy and soundfile are imported
lazily inside the functions that need them, which only run in worker environments).
"""

from __future__ import annotations

import importlib
import io
import math
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING

from nida.bench.protocol import WorkRequest
from nida.utils.io import atomic_write_bytes

if TYPE_CHECKING:
    import numpy.typing as npt

__all__ = [
    "REQUIRED_SAMPLE_RATE_HZ",
    "AudioInfo",
    "FrameCountError",
    "audio_info",
    "finish_frame_output",
    "out_path_param",
    "require_16k_mono",
    "verify_frame_count",
    "write_float32_npy",
]

#: Input rate the orchestrator guarantees for the 16 kHz workers (diarization, CTC, emotion).
REQUIRED_SAMPLE_RATE_HZ = 16000


class FrameCountError(ValueError):
    """A model returned a number of frames that does not match the audio duration."""


@dataclass(frozen=True)
class AudioInfo:
    """Header facts about one audio file.

    Attributes:
        sample_rate_hz: Sampling rate of the file.
        num_channels: Channel count.
        num_samples: Samples per channel.
        duration_ms: ``1000 * num_samples / sample_rate_hz``.
    """

    sample_rate_hz: int
    num_channels: int
    num_samples: int
    duration_ms: float


def audio_info(path: str) -> AudioInfo:
    """Read the header of ``path`` with soundfile (no decoding).

    Raises:
        ValueError: if the file reports a non-positive rate, channel count or length.
        RuntimeError: if soundfile cannot open the file (``soundfile.LibsndfileError``).
    """
    soundfile = importlib.import_module("soundfile")
    info = soundfile.info(path)
    rate_hz = int(info.samplerate)
    channels = int(info.channels)
    samples = int(info.frames)
    if rate_hz <= 0 or channels <= 0 or samples <= 0:
        raise ValueError(f"{path}: empty or invalid audio ({rate_hz} Hz, {channels} ch, {samples})")
    return AudioInfo(rate_hz, channels, samples, 1000.0 * samples / rate_hz)


def require_16k_mono(request: WorkRequest) -> AudioInfo:
    """Return the audio header of ``request`` after checking it is 16 kHz mono.

    The orchestrator is responsible for producing 16 kHz mono files; anything else is an
    error here rather than a silent resample (a wrong rate shifts every timestamp).

    Raises:
        ValueError: if the file is not 16 kHz mono.
    """
    info = audio_info(request.audio_path)
    if info.sample_rate_hz != REQUIRED_SAMPLE_RATE_HZ or info.num_channels != 1:
        raise ValueError(
            f"{request.clip_id}: expected {REQUIRED_SAMPLE_RATE_HZ} Hz mono, got "
            f"{info.sample_rate_hz} Hz, {info.num_channels} channel(s)"
        )
    return info


def out_path_param(request: WorkRequest) -> Path:
    """The ``.npy`` output path named by ``request.params["out_path"]``.

    Raises:
        ValueError: if the parameter is missing or does not end in ``.npy``.
    """
    raw = request.params.get("out_path")
    if not raw:
        raise ValueError(f"{request.clip_id}: missing params.out_path")
    path = Path(raw)
    if path.suffix != ".npy":
        raise ValueError(f"{request.clip_id}: params.out_path must end in .npy, got {raw!r}")
    return path


def verify_frame_count(
    *,
    num_frames: int,
    frame_hop_ms: float,
    duration_ms: float,
    tolerance_ms: float,
    what: str,
) -> None:
    """Check that ``num_frames`` frames of ``frame_hop_ms`` cover ``duration_ms``.

    A wrong hop assumption (say 20 ms instead of 10 ms) would silently shift every
    downstream timestamp, so each worker verifies the model output against the audio length.

    Raises:
        FrameCountError: if there are no frames, or ``|num_frames * hop - duration|`` exceeds
            ``tolerance_ms``.
        ValueError: on a non-positive hop or negative tolerance.
    """
    if frame_hop_ms <= 0 or tolerance_ms < 0:
        raise ValueError(f"invalid hop/tolerance: {frame_hop_ms} ms / {tolerance_ms} ms")
    if num_frames < 1:
        raise FrameCountError(f"{what}: model returned no frames for {duration_ms:.1f} ms of audio")
    covered_ms = num_frames * frame_hop_ms
    if not math.isfinite(covered_ms) or abs(covered_ms - duration_ms) > tolerance_ms:
        raise FrameCountError(
            f"{what}: {num_frames} frames x {frame_hop_ms} ms = {covered_ms:.1f} ms does not "
            f"match the audio duration {duration_ms:.1f} ms (tolerance {tolerance_ms} ms)"
        )


def write_float32_npy(path: Path, array: npt.ArrayLike) -> None:
    """Atomically write ``array`` as a float32 ``.npy`` file (no pickle).

    Raises:
        ValueError: if the array contains NaN or infinity (nothing is written).
    """
    import numpy as np

    data = np.ascontiguousarray(array, dtype=np.float32)
    if not np.isfinite(data).all():
        raise ValueError(f"{path}: refusing to write non-finite values")
    buffer = io.BytesIO()
    np.save(buffer, data, allow_pickle=False)
    atomic_write_bytes(path, buffer.getvalue())


def finish_frame_output(
    *,
    out_path: Path,
    array: npt.ArrayLike,
    frame_hop_ms: float,
    duration_ms: float,
    tolerance_ms: float,
    what: str,
) -> dict[str, object]:
    """Verify the frame count of a ``[num_frames, num_columns]`` array, write it, describe it.

    The check runs before the write, so a result with a wrong frame count never produces an
    output file.

    Returns:
        Payload fragment: ``out_path``, ``num_frames``, ``num_columns``, ``frame_hop_ms`` and
        ``duration_ms``.

    Raises:
        FrameCountError: on a frame count that does not fit the duration.
        ValueError: if the array is not 2-D or contains non-finite values.
    """
    import numpy as np

    data = np.asarray(array)
    if data.ndim != 2:
        raise ValueError(
            f"{what}: expected a 2-D [num_frames, num_columns] array, got {data.shape}"
        )
    num_frames, num_columns = int(data.shape[0]), int(data.shape[1])
    verify_frame_count(
        num_frames=num_frames,
        frame_hop_ms=frame_hop_ms,
        duration_ms=duration_ms,
        tolerance_ms=tolerance_ms,
        what=what,
    )
    write_float32_npy(out_path, data)
    return {
        "out_path": str(out_path),
        "num_frames": num_frames,
        "num_columns": num_columns,
        "frame_hop_ms": frame_hop_ms,
        "duration_ms": round(duration_ms, 3),
    }