inference/core/models/utils/batching.py

from typing import Generator, Iterable, List, TypeVar, Union

B = TypeVar("B")


def calculate_input_elements(input_value: Union[B, List[B]]) -> int:
    return len(input_value) if issubclass(type(input_value), list) else 1


def create_batches(
    sequence: Iterable[B], batch_size: int
) -> Generator[List[B], None, None]:
    batch_size = max(batch_size, 1)
    current_batch = []
    for element in sequence:
        if len(current_batch) == batch_size:
            yield current_batch
            current_batch = []
        current_batch.append(element)
    if len(current_batch) > 0:
        yield current_batch