Method batch_generate

mcts_math/react_batch.py:83–144

(self, questions: List[str])

        return solver
    
    def batch_generate(self, questions: List[str]):

        solvers = [REACT(config=self.config, question=question) for question in questions]

        for step in tqdm(range(self.config.max_depth), desc="Step Processing"):
            prompts = []
            epoch_solvers = []
            next_solvers = []

            for solver in solvers:
                if solver.should_generate_next():
                    prompts.append(solver.create_prompt())
                    epoch_solvers.append(solver)
                else:
                    next_solvers.append(solver)
            
            next_solver_span = len(next_solvers)
            if len(epoch_solvers) < 1:
                break
            
            # vllm run
            outputs = self.llm(prompts)
            # post-process outputs
            with ProcessPool(max_workers=min(len(epoch_solvers), os.cpu_count())) as pool:
                future = pool.map(self.__class__.processor, epoch_solvers, outputs, timeout=TIMEOUT_SECONDS)
                iterator = future.result()
            
            if len(epoch_solvers) > 100:  
                progress_bar = tqdm(total=len(epoch_solvers), desc="Execute")  
            else:  
                progress_bar = None 

            while True:
                try:
                    result = next(iterator)
                    next_solvers.append(result)
                except StopIteration:
                    break
                except TimeoutError as error:
                    next_solvers.append(None)
                    print(error)
                except Exception as error:
                    print(error)
                    next_solvers.append(None)
                if progress_bar is not None:
                    progress_bar.update(1) 
            
            if progress_bar is not None:
                progress_bar.close() 

            # update solvers
            assert len(epoch_solvers) == len(next_solvers[next_solver_span:]), f"Data is not matched, {len(epoch_solvers)} vs {len(next_solvers[next_solver_span:])}."
            for idx, (ori_solver, new_solver) in enumerate(zip(epoch_solvers, next_solvers[next_solver_span:])):
                if new_solver is None:
                    next_solvers[next_solver_span + idx] = ori_solver
            solvers = next_solvers

        jsonlines = {}