| """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", |
| ] |
|
|
| |
| 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), |
| } |
|
|
