AutoSurveys / AutoSurvey

134 stars 5 forks source link

a better choice of bacth_chat thread #8

Open denislov opened 1 month ago

denislov commented 1 month ago
def batch_chat(self, text_batch, temperature=0):
    res_l = ["No response"] * len(text_batch)
    with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
        futures = {
            executor.submit(self.__chat, text, temperature, res_l, i): i
            for i, text in enumerate(text_batch)
        }
        for future in tqdm(
            concurrent.futures.as_completed(futures), total=len(futures)
        ):
            idx = futures[future]
            try:
                future.result()
            except Exception as exc:
                print(f"Thread {idx} generated an exception: {exc}")
            else:
                print(f"Thread {idx} completed successfully")

    return res_l