| 42 | |
| 43 | |
| 44 | async def get_debug_models(config: Config) -> Tuple["ModelWrapper", List["ModelWrapper"]]: |
| 45 | from trinity.common.models.allocator import Allocator |
| 46 | |
| 47 | allocator = Allocator(config.explorer) |
| 48 | rollout_model = allocator.get_model(config.explorer.rollout_model, "rollout", 0) |
| 49 | auxiliary_models = [ |
| 50 | allocator.get_model(auxiliary_model_config, f"auxiliary_{index}", 0) |
| 51 | for index, auxiliary_model_config in enumerate(config.explorer.auxiliary_models) |
| 52 | ] |
| 53 | await asyncio.gather( |
| 54 | rollout_model.prepare(), |
| 55 | *[auxiliary_model.prepare() for auxiliary_model in auxiliary_models], |
| 56 | ) |
| 57 | return rollout_model, auxiliary_models |
| 58 | |
| 59 | |
| 60 | def debug( |