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