espnet3.systems.base.inference_runner.InferenceRunner

source

class espnet3.systems.base.inference_runner.InferenceRunner(provider: EnvironmentProvider, idx_key: str = 'utt_id', hyp_key: str | Sequence[str] = 'hyp', ref_key: str | Sequence[str] = 'ref', **kwargs)

Bases: BaseRunner

Execute a configured backend and persist its formatted results.

input_key, inference contract and output_artifacts belong to the recipe. The runner checks inputs and outputs when a contract is configured. Legacy configs may still provide output_fn.

Initialize the inference runner with output key settings.

• Parameters:
  • provider – Environment provider that supplies dataset/model/env.
  • idx_key – Output dict key used as the sample identifier written in the first column of each SCP line. This ties each inference result back to its dataset sample. Defaults to "utt_id".
  • hyp_key – Hypothesis key or keys expected in the output dict.
  • ref_key – Reference key or keys expected in the output dict.
  • **kwargs – Forwarded to BaseRunner.

static close_writers(writers: Dict[str, Any], state: Dict[str, Any], **env) → Dict[str, Any] | None

Close shard-local SCP files and report which output keys were written.

static format_output(result, , batched)

Convert raw model results to named outputs before contract validation.

• Parameters:
  • result – What the backend returned.
  • batched – Whether the call was a batch.
• Returns: The named outputs; result unless a system overrides it.

######### Examples

>>> InferenceRunner.format_output({"text": "hi"}, batched=False)
{'text': 'hi'}

classmethod forward(idx, dataset=None, model=None, **kwargs)

Call model with one item or a batch from dataset.

Example conf/inference.yaml:

input_key: speech
model:
  _target_: src.inference.MyBackend
contract:
  inputs:
    speech: {kind: audio, sample_rate: 16000,
             rate_policy: resample, label: Input Audio}
  outputs:
    text: {kind: text, normalize: [strip, to_lower],
           label: Transcription}

forward checks inputs, calls the model, formats its output, and validates the result. Both infer and InferenceModel use this path:

dataset[0] = {"utt_id": "u1", "speech": (44100, waveform)}
# forward receives Audio at 16 kHz with every channel preserved.
model(speech=audio_16k)  # -> {"text": "  HELLO  "}
# boundary returns {"text": "hello", "utt_id": "u1"}

input_key must match contract.inputs. Outputs must have the declared names; {"hyp": ...} above raises an error. A bare audio array has no rate to check; pass Audio, (rate, array), or a path when the rate matters. output_fn and contract cannot coexist.

Override format_output to adapt raw model results:

class ASRRunner(InferenceRunner):
    @staticmethod
    def format_output(result, *, batched):
        def best(hypotheses):
            return {"text": hypotheses[0][0] or ""}
        return [best(r) for r in result] if batched else best(result)

If overriding forward itself, apply the shared checks explicitly:

@staticmethod
def forward(idx, dataset=None, model=None, **kwargs):
    contract = resolve_inference_contract(
        kwargs.get("inference_contract"),
        kwargs.get("inference_contract_root"),
    )
    if contract is None:
        return model(speech=dataset[idx]["speech"])
    contract.validate_input_keys(kwargs["input_key"])
    prepared, original = prepare_inference_inputs(
        dataset, idx, contract, model
    )
    result = model(speech=prepared[idx]["speech"])
    return validate_inference_outputs(
        result, idx, contract, model, original,
        kwargs.get("idx_key", "utt_id"),
    )

Optional inputs must also appear in input_key:

input_key: [speech, prompt]
contract:
  inputs:
    speech: {kind: audio, sample_rate: 16000}
    prompt: {kind: text, optional: true}
  outputs:
    text: {kind: text}
# Missing prompt -> model(speech=...); present ->
# model(speech=..., prompt=...). Do not mix in one batch.

Custom normalizer example:

# inference.yaml: contract.imports: [src.contracts]
# contract.outputs.text.normalize: [strip, my_normalizer]
# src/contracts.py:
from espnet3.contract import register_normalizer
register_normalizer("my_normalizer", normalize_text)

Custom forward implementations own their contract checks; calling super().forward uses the default checked path. format_output receives raw model results and returns named outputs before validation. Changes after super().forward require output validation again. Bundled recipe code needs trust_user_code=True when loaded. Without a contract, the legacy output_fn path still applies.

• Parameters:
  • idx – One dataset index, or a list of them for a batch.
  • dataset – The dataset to read items from.
  • model – The backend called with the checked inputs.
  • **kwargs – input_key, inference_contract, idx_key, model_kwargs and the legacy output_fn / output_fn_path.
• Returns: One named output mapping, or a list of them for a batch.

######### Examples

>>> from espnet3.contract.inference import load_inference_contract
>>> contract = load_inference_contract(
...     {
...         "inputs": {"text": {"kind": "text"}},
...         "outputs": {"text": {"kind": "text"}},
...     }
... )
>>> InferenceRunner.forward(
...     0,
...     dataset=[{"text": "hi"}],
...     model=lambda text: {"text": text},
...     input_key="text",
...     inference_contract=contract,
... )
{'text': 'hi', 'utt_id': '0'}

merge(shard_dirs: List[Path]) → None

Merge per-shard SCP files into the test-set output directory.

Reads field_keys.txt from each shard to discover output field names, then concatenates each <field>.scp across shards in shard order into output_dir / shard_subdir.

• Parameters:shard_dirs – Completed shard directories in shard-id order.
• Raises:RuntimeError – If no output keys are found across all shards.

static open_writers(shard_dir: Path | None, output_artifacts: Dict[str, dict] | None = None, **env) → Dict[str, Any]

Open per-shard SCP writers for worker-side inference outputs.

static prepare_model_inputs(inputs)

Adapt checked Kind values to a backend’s argument representation.

• Parameters:inputs – The backend’s keyword arguments, already checked by their kinds (a list per key for a batch).
• Returns: The arguments the backend is called with; inputs unless a system overrides it.

######### Examples

>>> InferenceRunner.prepare_model_inputs({"text": "hi"})
{'text': 'hi'}

resolve_idx_key(output: Dict[str, Any]) → str

Validate that the configured sample-identifier key exists in output.

static write_record(writers: Dict[str, Any], result: Any, state: Dict[str, Any], idx_key: str = 'utt_id', output_keys=None, hyp_key=None, ref_key=None, **env) → None

Validate one forward result and stream it into shard-local SCP files.