| 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 |
|
|
