| 8 | from concurrent.futures import ThreadPoolExecutor, as_completed |
| 9 | |
| 10 | def execute(func, input_list_or_num_samples: Union[list, int], output_path: str, max_workers: int, logger: logging.Logger): |
| 11 | out = open(output_path, 'a') |
| 12 | with ThreadPoolExecutor(max_workers=max_workers) as executor: |
| 13 | if isinstance(input_list_or_num_samples, list): |
| 14 | futures = [executor.submit(func, item) for item in input_list_or_num_samples] |
| 15 | else: |
| 16 | futures = [executor.submit(func) for _ in range(input_list_or_num_samples)] |
| 17 | for future in tqdm(as_completed(futures), total=len(futures)): |
| 18 | try: |
| 19 | res: dict = future.result() |
| 20 | except Exception as e: |
| 21 | logger.info(f"[error] {e}") |
| 22 | continue |
| 23 | if res: |
| 24 | out.write(json.dumps(res, ensure_ascii=False) + '\n') |
| 25 | out.flush() |
| 26 | out.close() |
| 27 | |
| 28 | def retry(max: int=10, sleep: int=1, logger: logging.Logger=None): |
| 29 | def decorator(func): |