Add model_name for auxiliary models (#461)

Co-authored-by: chenyushuo <[email protected]>

Author: hiyuchang

.pre-commit-config.yaml

@@ -15,7 +15,6 @@ repos:
     rev: 23.7.0
     hooks:
     - id: black
-      language_version: python3.12
       args: [--line-length=100]
 
   - repo: https://github.com/pycqa/isort

trinity/common/config.py

@@ -490,6 +490,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

@@ -70,6 +70,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."""
@@ -113,6 +117,7 @@ def __init__(
             engine_type.startswith("vllm") or engine_type == "tinker"
         ), "Only vLLM and tinker model is supported for now."
         self.model = model
+        self._model_name = None
         self.api_address: str = None
         self.openai_client: openai.OpenAI = None
         self.openai_async_client: openai.AsyncOpenAI = None
@@ -128,6 +133,7 @@ def __init__(
 
     async def prepare(self) -> None:
         """Prepare the model wrapper."""
+        self._model_name = await self.model.get_model_name.remote()
         self.api_address = await self.model.get_api_server_url.remote()
         if self.api_address is None:
             self.logger.info("API server is not enabled for inference model.")
@@ -285,6 +291,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 self._model_name
+
+    @property
+    async def model_name_async(self) -> Optional[str]:
+        """Get the name of the model."""
+        return self._model_name
+
     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

@@ -718,6 +718,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.config.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,