Method create_llm

mcts_math/agents/react.py:69–93

(self)

        )

    def create_llm(self) -> Callable[[...], List[str]]:
        GPUS = os.environ.get('CUDA_VISIBLE_DEVICES', "0").split(',')
        llm = LLM(
            model=self.config.model_dir, 
            tensor_parallel_size=len(GPUS), 
            trust_remote_code=True,
            seed=self.config.seed,
            swap_space=self.config.swap_space,
        )
        sampling_params = SamplingParams(
            top_k=self.config.top_k,
            top_p=self.config.top_p,
            use_beam_search=self.config.use_beam_search,
            best_of=self.config.best_of,
            max_tokens=self.config.max_tokens, 
            stop=self.stop,
            #seed=self.config.seed,
        )
        return partial(
            local_vllm,
            llm=llm,
            sampling_params=sampling_params,
            n=1,
            temperature=self.config.temperature,
        )

    def should_generate_next(self) -> bool:
        return not self.current_node.is_terminal and self.current_node.depth <= self.config.max_depth