| 26 | |
| 27 | |
| 28 | class REACT(BaseTree): |
| 29 | |
| 30 | REACT_NODE_KEYS: List[str] = ["action", "action_input", "final_answer"] |
| 31 | prompt_wrap: Optional[Callable[[...], str]] = None |
| 32 | obs_wrap: Optional[Callable[str, str]] = None |
| 33 | step_unwrap: Optional[Callable[[...], Dict[str, str]]] = None |
| 34 | |
| 35 | def __init__(self, **kwargs) -> None: |
| 36 | super().__init__(**kwargs) |
| 37 | |
| 38 | if self.config.prompt_wrap == "react": |
| 39 | from .utils import react_prompt_wrap, react_obs_wrap, react_step_result_unwrap |
| 40 | |
| 41 | self.prompt_wrap = react_prompt_wrap |
| 42 | self.obs_wrap = react_obs_wrap |
| 43 | self.step_unwrap = react_step_result_unwrap |
| 44 | |
| 45 | elif self.config.prompt_wrap == "react_sft": |
| 46 | from .utils import react_sft_prompt_wrap, react_sft_obs_wrap, react_sft_step_result_unwrap |
| 47 | |
| 48 | self.prompt_wrap = react_sft_prompt_wrap |
| 49 | self.obs_wrap = react_sft_obs_wrap |
| 50 | self.step_unwrap = react_sft_step_result_unwrap |
| 51 | |
| 52 | @field_validator("config") |
| 53 | def validate_config(cls, cfg: Any): |
| 54 | super().validate_config(cfg) |
| 55 | if not cfg.mode == "react": |
| 56 | raise ValueError(f"Wrong value for config mode, must be react") |
| 57 | if not cfg.n_generate_sample == 1: |
| 58 | raise ValueError(f"Wrong value for config n_generate_sample, must be 1") |
| 59 | if cfg.stop is None: |
| 60 | raise ValueError(f"Wrong value for config stop, cannot be None") |
| 61 | return cfg |
| 62 | |
| 63 | def create_node(self, parent: Optional[Type[BaseNode]] = None) -> Type[BaseNode]: |
| 64 | return BaseNode( |
| 65 | parent=parent, |
| 66 | additional_state_keys=self.REACT_NODE_KEYS, |
| 67 | ) |
| 68 | |
| 69 | def create_llm(self) -> Callable[[...], List[str]]: |
| 70 | GPUS = os.environ.get('CUDA_VISIBLE_DEVICES', "0").split(',') |
| 71 | llm = LLM( |
| 72 | model=self.config.model_dir, |
| 73 | tensor_parallel_size=len(GPUS), |
| 74 | trust_remote_code=True, |
| 75 | seed=self.config.seed, |
| 76 | swap_space=self.config.swap_space, |
| 77 | ) |
| 78 | sampling_params = SamplingParams( |
| 79 | top_k=self.config.top_k, |
| 80 | top_p=self.config.top_p, |
| 81 | use_beam_search=self.config.use_beam_search, |
| 82 | best_of=self.config.best_of, |
| 83 | max_tokens=self.config.max_tokens, |
| 84 | stop=self.stop, |
| 85 | #seed=self.config.seed, |
no outgoing calls
no test coverage detected