nemo_curator.models.asr.base

Stage-adapter contract for audio speech-recognition.

ASRStage owns Curator-side glue (task.data reads, batching, ISO language mapping, _skipme), while ASRAdapter owns the model-side call (weight download, model loading, generation, and packing into ASRResult). The split lets the stage swap models via a single YAML adapter_target: line.

Module Contents

Classes

Name	Description
`ASRAdapter`	Structural protocol every ASR adapter must implement.
`ASRResult`	Canonical per-utterance ASR adapter output.

API

class nemo_curator.models.asr.base.ASRAdapter()

Protocol

Structural protocol every ASR adapter must implement.

ASRStage constructs adapters with model_id and the explicitly configured adapter_kwargs. Model-provider options therefore stay with the adapter that implements them instead of becoming shared stage fields.

Per-batch contract: transcribe_batch receives a list of per-task dicts (unpacked from task.data) and returns one ASRResult per input, in order. Expected per-item keys (stage-populated):

waveform: contiguous, mono, 1-D float32 NumPy samples normalized by ASRStage from a file or a reader-provided in-memory waveform.
sample_rate (int): the stage’s configured target sample rate.
audio_seconds (float): duration of this model-input chunk.
language (str | None): human-readable name (e.g. "English").
language_code (str | None): original language code from the configured stage input column.
task_id (str | None): carried through for diagnostics.

model_id

str

Identifier of the underlying model checkpoint.

nemo_curator.models.asr.base.ASRAdapter.download_weights_on_node() -> None

Download weights to local cache without allocating a GPU.

The stage calls this once per node on a lightweight adapter instance so provider-specific download options remain encapsulated by that adapter.

nemo_curator.models.asr.base.ASRAdapter.load_model(
    num_gpus: int
) -> None

Load the model into the worker process using the stage-owned GPU count.

ASRStage derives num_gpus from its Curator resource request so adapters do not expose a second, independently configurable GPU count.

nemo_curator.models.asr.base.ASRAdapter.transcribe_batch(
    items: list[dict[str, typing.Any]]
) -> list[nemo_curator.models.asr.base.ASRResult]

Run inference on a batch of per-task dicts.

Returns one ASRResult per input, in order; skipped items must still appear with skipped=True to preserve task ordering.

nemo_curator.models.asr.base.ASRAdapter.unload_model() -> None

Release GPU memory and worker-local state.

class nemo_curator.models.asr.base.ASRResult(
    text: str,
    skipped: bool = False,
    skip_reason: str | None = None,
    unsupported_language: str | None = None,
    extras: dict[str, typing.Any] = dict()
)

Dataclass

Canonical per-utterance ASR adapter output.

Identical across every adapter so the stage’s schema-mutation path stays constant when the adapter is swapped.

extras

dict[str, Any] = field(default_factory=dict)

Adapter-specific, manifest-serializable diagnostics. When the stage’s extras_key is enabled, it writes a shallow copy of this dictionary under that one nested output field without interpreting individual keys.

skip_reason

str | None = None

Optional machine-readable reason written to _skipme when skipped is true. Defaults to "empty_audio" in the stage.

skipped

bool = False

True when the item could not be processed. The stage writes skip_reason to _skipme and falls back to "empty_audio" when no reason is supplied.

text

str

Transcription text. Empty if skipped.

unsupported_language

str | None = None

Optional normalized language code used by the stage to annotate items excluded by its language allowlist.