Add model_name for auxiliary models

trinity/common/config.py

@@ -477,6 +477,7 @@ class ModelConfig:
 class InferenceModelConfig:
     # ! DO NOT SET in explorer.rollout_model, automatically set from config.model.model_path
     model_path: Optional[str] = None
+    name: Optional[str] = None
 
     engine_type: str = "vllm"
     engine_num: int = 1

trinity/common/models/model.py

@@ -66,6 +66,10 @@ def get_model_path(self) -> Optional[str]:
         """Get the model path"""
         return None
 
+    def get_model_name(self) -> Optional[str]:
+        """Get the name of the model."""
+        return None
+
 
 def _history_recorder(func):
     """Decorator to record history of the model calls."""
@@ -279,6 +283,16 @@ async def model_path_async(self) -> str:
         """Get the model path."""
         return await self.model.get_model_path.remote()
 
+    @property
+    def model_name(self) -> Optional[str]:
+        """Get the name of the model."""
+        return ray.get(self.model.get_model_name.remote())
+
+    @property
+    async def model_name_async(self) -> Optional[str]:
+        """Get the name of the model."""
+        return await self.model.get_model_name.remote()
+
     def get_lora_request(self) -> Any:
         if self.enable_lora:
             return ray.get(self.model.get_lora_request.remote())

trinity/common/models/vllm_model.py

@@ -39,6 +39,7 @@ def __init__(
         import vllm
         from vllm.sampling_params import RequestOutputKind
 
+        self.name = config.name
         self.logger = get_logger(__name__)
         self.vllm_version = get_vllm_version()
         self.config = config
@@ -718,6 +719,9 @@ def get_model_version(self) -> int:
     def get_model_path(self) -> str:
         return self.config.model_path  # type: ignore [return-value]
 
+    def get_model_name(self) -> Optional[str]:
+        return self.name  # type: ignore [return-value]
+
     def get_lora_request(self, lora_path: Optional[str] = None) -> Any:
         from vllm.lora.request import LoRARequest
 

trinity/common/workflows/agentscope_workflow.py

@@ -118,13 +118,18 @@ def __init__(
                 "top_logprobs": self.task.rollout_args.logprobs,
             },
         )
-        self.auxiliary_chat_models = [
-            TrinityChatModel(
-                openai_async_client=aux_model,
-                # TODO: customize generate_kwargs for auxiliary models if needed
-            )
-            for aux_model in (self.auxiliary_models or [])
-        ]
+
+        # TODO: customize generate_kwargs for auxiliary models if needed
+        if self.auxiliary_model_wrappers is not None and self.auxiliary_models is not None:
+            self.auxiliary_chat_models = {
+                aux_model_wrapper.model_name
+                or f"auxiliary_model_{i}": TrinityChatModel(openai_async_client=aux_model)
+                for i, (aux_model_wrapper, aux_model) in enumerate(
+                    zip(self.auxiliary_model_wrappers, self.auxiliary_models)
+                )
+            }
+        else:
+            self.auxiliary_chat_models = {}
 
     def construct_experiences(
         self,