| 46 | |
| 47 | |
| 48 | class Solver(BaseModel): |
| 49 | model_config = ConfigDict(arbitrary_types_allowed=True) |
| 50 | |
| 51 | config: Any |
| 52 | |
| 53 | stop: List[str] = None |
| 54 | |
| 55 | llm: Optional[Callable[[...], List[str]]] = None |
| 56 | |
| 57 | llm_model_id: Optional[str] = None |
| 58 | engine: Optional[LLM] = None |
| 59 | generate_sampling_params: Optional[SamplingParams] = None |
| 60 | value_sampling_params: Optional[SamplingParams] = None |
| 61 | need_value_func: bool = False |
| 62 | max_solver_steps: int = 1 |
| 63 | |
| 64 | def __init__(self, **kwargs) -> None: |
| 65 | super().__init__(**kwargs) |
| 66 | |
| 67 | self.llm_model_id = self.config.model_dir |
| 68 | |
| 69 | if self.config.stop: |
| 70 | # omegaconf.listconfig.ListConfig -> list |
| 71 | self.stop = OmegaConf.to_object(self.config.stop) |
| 72 | |
| 73 | self.llm = self.create_llm() |
| 74 | self.need_value_func = self.config.need_value_func |
| 75 | |
| 76 | if self.config.mode == "sbs": |
| 77 | self.max_solver_steps = self.config.max_depth |
| 78 | elif self.config.mode == "mcts": |
| 79 | self.max_solver_steps = self.config.iterations |
| 80 | self.config.step_beam_width = 1 |
| 81 | |
| 82 | @field_validator("config") |
| 83 | def validate_config(cls, cfg: Any): |
| 84 | if issubclass(type(cfg), DictConfig): |
| 85 | return cfg |
| 86 | |
| 87 | raise TypeError("Wrong type for `config`, must be subclass of BaseConfig") |
| 88 | |
| 89 | def create_llm(self) -> Callable[[...], List[str]]: |
| 90 | if self.config.seed: |
| 91 | set_seed(self.config.seed) |
| 92 | engine, sampling_params = llm_engine(self.config) |
| 93 | self.engine = engine |
| 94 | self.generate_sampling_params = sampling_params |
| 95 | self.value_sampling_params = copy.deepcopy(sampling_params) |
| 96 | self.value_sampling_params.max_tokens = 1 |
| 97 | self.value_sampling_params.n = 1 |
| 98 | return partial( |
| 99 | local_generator, |
| 100 | engine=self.engine, |
| 101 | ) |
| 102 | |
| 103 | @staticmethod |
| 104 | def processor(solver: BaseTree, output: List[RequestOutput]) -> BaseTree: |
| 105 | solver.generate_next_step(output) |