diff --git a/great_ai/utilities/__init__.py b/great_ai/utilities/__init__.py index 5bc97c4..65c2a17 100644 --- a/great_ai/utilities/__init__.py +++ b/great_ai/utilities/__init__.py @@ -5,6 +5,11 @@ from .evaluate_ranking import evaluate_ranking from .get_sentences import get_sentences from .language import english_name_of_language, is_english, predict_language from .logger import get_logger -from .parallel_map import WorkerException, parallel_map, threaded_parallel_map +from .parallel_map import ( + WorkerException, + parallel_map, + simple_parallel_map, + threaded_parallel_map, +) from .unchunk import unchunk from .unique import unique diff --git a/great_ai/utilities/parallel_map/__init__.py b/great_ai/utilities/parallel_map/__init__.py index 7b6ceec..a984057 100644 --- a/great_ai/utilities/parallel_map/__init__.py +++ b/great_ai/utilities/parallel_map/__init__.py @@ -1,3 +1,4 @@ from .parallel_map import parallel_map +from .simple_parallel_map import simple_parallel_map from .threaded_parallel_map import threaded_parallel_map from .worker_exception import WorkerException diff --git a/great_ai/utilities/parallel_map/simple_parallel_map.py b/great_ai/utilities/parallel_map/simple_parallel_map.py new file mode 100644 index 0000000..3dca6fd --- /dev/null +++ b/great_ai/utilities/parallel_map/simple_parallel_map.py @@ -0,0 +1,30 @@ +from typing import Awaitable, Callable, List, Optional, Sequence, TypeVar, Union + +from tqdm.cli import tqdm + +from .parallel_map import parallel_map + +T = TypeVar("T") +V = TypeVar("V") + + +def simple_parallel_map( + func: Callable[[T], Union[V, Awaitable[V]]], + input_values: Sequence[T], + *, + chunk_size: Optional[int] = None, + concurrency: Optional[int] = None, +) -> List[V]: + input_values = list(input_values) # in case the input is mistakenly not a sequence + return list( + tqdm( + parallel_map( + func=func, + input_values=input_values, + chunk_size=chunk_size, + concurrency=concurrency, + ), + total=len(input_values), + dynamic_ncols=True, + ) + )