kaira.models.generic.ParallelModel

Inheritance diagram for ParallelModel

class kaira.models.generic.ParallelModel(max_workers: int | None = None, steps: List[Tuple[str, Callable]] | None = None, branches: List[Callable] | None = None, aggregator: Callable | None = None)

Bases: ConfigurableModel

A model that processes steps in parallel.

All steps receive the same input data and process independently.

Methods

__init__	Initialize the parallel model.
add_step	Add a processing step to the model with an optional name.
forward	Execute the model in parallel on the input data.
from_config	Create model instance from configuration.
from_hydra_config	Create model from Hydra DictConfig.
from_pretrained_config	Create model from Hugging Face PretrainedConfig.
remove_step	Remove a processing step from the model.

Attributes

step_configs

__init__(max_workers: int | None = None, steps: List[Tuple[str, Callable]] | None = None, branches: List[Callable] | None = None, aggregator: Callable | None = None)

Initialize the parallel model.

Parameters:

• max_workers – Maximum number of worker threads (None uses default ThreadPoolExecutor behavior)

• steps – Optional initial list of named processing steps as (name, step) tuples

• branches – Alternative way to specify processing steps as a list of callables

• aggregator – Optional function to aggregate results (if None, returns dict of outputs)

step_configs: List[Tuple[str, Callable]]

add_step(step: Callable, name: str | None = None)

Add a processing step to the model with an optional name.

Parameters:

• step – A callable function or object that processes input data

• name – Optional name for the step (auto-generated if None)

Returns:

The model instance for method chaining

Raises:

TypeError – If step is not callable

classmethod from_config(config, **kwargs)

Create model instance from configuration.

Parameters:

• config – Configuration object (PretrainedConfig, DictConfig, or dict)

• **kwargs – Additional parameters to override config

Returns:

Model instance

classmethod from_hydra_config(config: DictConfig, **kwargs)

Create model from Hydra DictConfig.

Parameters:

• config – Hydra configuration

• **kwargs – Additional parameters

Returns:

Model instance

classmethod from_pretrained_config(config: PretrainedConfig, **kwargs)

Create model from Hugging Face PretrainedConfig.

Parameters:

• config – PretrainedConfig instance

• **kwargs – Additional parameters

Returns:

Model instance

remove_step(index: int)

Remove a processing step from the model.

Parameters:

index – The index of the step to remove

Returns:

The model instance for method chaining

Raises:

IndexError – If the index is out of range

forward(input_data: Any, *args: Any, **kwargs: Any) → Any

Execute the model in parallel on the input data.

Parameters:

• input_data – The data to process

• *args – Additional positional arguments passed to each step.

• **kwargs – Additional keyword arguments passed to each step.

Returns:

Dictionary mapping step names to their respective outputs or aggregated results if an aggregator is provided