add use base format

Author: garyzhang99

trinity/common/config.py

@@ -49,6 +49,8 @@ class FormatConfig:
     # for unpaired preference dataset
     label_key: str = ""
 
+    use_base_format: bool = False
+
 
 @dataclass
 class GenerationConfig:

trinity/common/models/model.py

@@ -99,6 +99,12 @@ def convert_messages_to_experience(self, messages: List[dict]) -> Experience:
         else:
             return ray.get(self.model.convert_messages_to_experience.remote(messages))
 
+    def tokenize_text(self, text: str) -> Tensor:
+        if self.use_async:
+            return ray.get(self.model.tokenize_text_async.remote(text))
+        else:
+            return ray.get(self.model.tokenize_text.remote(text))
+
     def get_ckp_version(self) -> int:
         return ray.get(self.model.get_ckp_version.remote())
 

trinity/common/models/vllm_async_model.py

@@ -227,6 +227,24 @@ async def convert_messages_to_experience_async(self, messages: List[dict]) -> Ex
             action_mask=action_mask,
         )
 
+    async def tokenize_text_async(self, text: str) -> torch.Tensor:
+        """Convert text to token ids tensor.
+
+        Args:
+            text (str): Input text to be tokenized
+
+        Returns:
+            torch.Tensor: Token ids tensor
+        """
+        if self.tokenizer is None:
+            self.tokenizer = await self.async_llm.get_tokenizer()
+
+        # Tokenize the text
+        token_ids = self.tokenizer.encode(text)
+
+        # Convert to tensor
+        return torch.tensor(token_ids, dtype=torch.int32)
+
     def shutdown(self):
         """Shutdown the vLLM v1 engine. This kills child processes forked
         by the vLLM engine. If not called, the child processes will be

trinity/common/models/vllm_model.py

@@ -268,6 +268,19 @@ def convert_messages_to_experience(self, messages: List[dict]) -> Experience:
             action_mask=action_mask,
         )
 
+    def tokenize_text(self, text: str) -> torch.Tensor:
+        """
+        Shape text to token ids.
+
+        Args:
+            text (str)
+
+        Returns:
+            torch.Tensor: token ids tensor
+        """
+        token_ids = self.tokenizer.encode(text)
+        return torch.tensor(token_ids, dtype=torch.int32)
+
     def has_api_server(self) -> bool:
         return False
 

trinity/common/workflows/workflow.py

@@ -190,12 +190,27 @@ def format_messages(self):
             messages.append({"role": "assistant", "content": self.reply_prefix})
         return messages
 
+    def format_prompt(self):
+        prompt_text = ""
+        if self.system_prompt:
+            prompt_text += self.system_prompt
+            prompt_text += "\nTask:\n" + self.task_desc + "\nResponse:\n"
+        else:
+            prompt_text += "\nTask:\n" + self.task_desc + "\nResponse:\n"
+        return prompt_text
+
     def run(self) -> List[Experience]:
         # TODO: Optimize the generate function
-        messages = self.format_messages()
+        if self.format_args.use_base_format:
+            prompt_text = self.format_prompt()
+        else:
+            messages = self.format_messages()
 
         logger.debug("start chat")
-        responses = self.model.chat(messages, **self.rollout_args)
+        if self.format_args.use_base_format:
+            responses = self.model.generate([prompt_text], **self.rollout_args)
+        else:
+            responses = self.model.chat(messages, **self.rollout_args)
         for response in responses:
             reward = self.reward_fn(  # type: ignore [misc]
                 response=response.response_text,  # type: ignore [arg-type]