espnet3.systems.base.inference_provider.InferenceProvider

source

class espnet3.systems.base.inference_provider.InferenceProvider(inference_config: DictConfig = None, params: Dict[str, Any] | None = None)

Bases: EnvironmentProvider, ABC

EnvironmentProvider specialized for dataset/model inference setup.

This implementation focuses on constructing just the dataset and model and returning them as environment entries. It is suitable for both local and distributed execution.

Design: : - build_* helpers are defined as @staticmethod/class methods so they are easily serializable and reusable on workers.

• The worker setup function must not capture self to remain pickle-safe for Dask.

• Parameters:
  • inference_config (DictConfig) – Hydra configuration used to build dataset/model.
  • params (Dict *[*str , Any ] | None) – Optional additional key-value pairs that will be merged into the returned environment (e.g., device, tokenizer, beam size).

Notes

• Subclasses must implement build_dataset and build_model.
• self.inference_config.update(self.params) allows lightweight overrides (e.g., runtime overrides) but avoid mutating deep structures unless intended.

Initialize InferenceProvider object.

static build_dataset(config: DictConfig)

Construct and return the dataset instance.

Build a DataOrganizer from config.dataset and return the selected test set. During parallel or distributed execution, the config object passed here is the configuration that the user passed when instantiating the class.

• Parameters:config (DictConfig) – Configuration object that includes a plain dataset organizer config and test_set.
• Returns: Dataset object resolved from organizer.test[test_set].
• Return type: Any
• Raises:NotImplementedError – Always in the base class; implement in subclass.

Example

>>> # Minimal sketch; actual keys depend on your subclass
>>> from omegaconf import OmegaConf
>>> config = OmegaConf.create(
...     {
...         "dataset": {
...             "test": [
...                 {
...                     "name": "test",
...                     "data_src_args": {"split": "test"},
...                 }
...             ]
...         },
...     }
... )
>>> ds = MyInferenceProvider.build_dataset(config)

Notes

• Keep dataset initialization lightweight by using lazy loading or memory mapping when possible.
• Rely on fields already present in config instead of reading global state whenever possible.

build_env_local() → Dict[str, Any]

Build the environment once on the driver for local inference.

• Returns: Environment dict with at least: : > - "dataset": The instantiated dataset.

  • "model": The instantiated model. <br/> Any additional fields from params are also included.

• Return type: Dict[str, Any]

Example

>>> provider = InferenceProvider(config, params={"device": "cuda"})
>>> env = provider.build_env_local()
>>> env.keys()
dict_keys(["dataset", "model", "device"])

static build_model(config: DictConfig)

Build config.model directly on this worker’s logical device.

• Parameters:config – Inference config with a model entry. A recipe_dir entry makes the model build inside that directory, so relative paths in model resolve against it.
• Returns: The instantiated model.

Example

s

>>> from omegaconf import OmegaConf
>>> config = OmegaConf.create(
...     {"device": "cpu", "model": {"_target_": "torch.nn.Identity"}}
... )
>>> InferenceProvider.build_model(config)
Identity()

build_worker_setup_fn() → Callable[[], Dict[str, Any]]

Return a Dask worker setup function that builds dataset/model.

The returned function is executed once per worker process and must not capture self. It closes over the immutable config and params snapshot, then constructs the environment on each worker.

• Returns: A zero-argument setup function that returns {"dataset": ..., "model": ..., **params}.
• Return type: Callable[[], Dict[str, Any]]

Example

>>> provider = InferenceProvider(config, params={"device": "cuda:0"})
>>> setup_fn = provider.build_worker_setup_fn()
>>> env = setup_fn()
>>> "dataset" in env and "model" in env
True