| 23 | |
| 24 | |
| 25 | class REACTBatch(BaseModel): |
| 26 | |
| 27 | config: Any |
| 28 | |
| 29 | stop: List[str] = None |
| 30 | |
| 31 | llm: Callable[[...], List[RequestOutput]] = None |
| 32 | |
| 33 | def __init__(self, **kwargs) -> None: |
| 34 | super().__init__(**kwargs) |
| 35 | |
| 36 | if self.config.stop: |
| 37 | self.stop = OmegaConf.to_object(self.config.stop) |
| 38 | |
| 39 | self.llm = self.create_llm() |
| 40 | |
| 41 | @field_validator("config") |
| 42 | def validate_config(cls, cfg: Any): |
| 43 | if issubclass(type(cfg), DictConfig): |
| 44 | return cfg |
| 45 | |
| 46 | raise TypeError("Wrong type for `config`, must be subclass of BaseConfig") |
| 47 | |
| 48 | def create_llm(self): |
| 49 | GPUS = os.environ.get('CUDA_VISIBLE_DEVICES', "0").split(',') |
| 50 | llm = LLM( |
| 51 | model=self.config.model_dir, |
| 52 | tensor_parallel_size=len(GPUS), |
| 53 | trust_remote_code=True, |
| 54 | seed=self.config.seed, |
| 55 | swap_space=self.config.swap_space, |
| 56 | ) |
| 57 | sampling_params = SamplingParams( |
| 58 | temperature=self.config.temperature, |
| 59 | top_k=self.config.top_k, |
| 60 | top_p=self.config.top_p, |
| 61 | use_beam_search=self.config.use_beam_search, |
| 62 | best_of=self.config.best_of, |
| 63 | max_tokens=self.config.max_tokens, |
| 64 | n=1, |
| 65 | stop=self.stop, |
| 66 | #seed=self.config.seed, |
| 67 | ) |
| 68 | return partial( |
| 69 | llm.generate, |
| 70 | sampling_params=sampling_params, |
| 71 | ) |
| 72 | |
| 73 | @staticmethod |
| 74 | def processor(solver: REACT, output: RequestOutput) -> REACT: |
| 75 | step_result = output.outputs[0].text.strip() |
| 76 | try: |
| 77 | step_result, parser_result = solver.step_unwrap(step_result) |
| 78 | except Exception as e: |
| 79 | parser_result = None |
| 80 | solver.update_current_node(step_result, parser_result) |
| 81 | return solver |
| 82 | |