[BUGFIX]Fix tokenizer bug when getting action masks with enable_thinking arguments (agentscope-ai#316)

Co-authored-by: 问昊 <[email protected]>

Author: garyzhang99

trinity/common/models/utils.py

@@ -16,6 +16,7 @@ def tokenize_and_mask_messages_hf(
     messages: List[dict],
     tools: Optional[List[dict]] = None,
     chat_template: Optional[str] = None,
+    enable_thinking: Optional[bool] = None,
 ) -> Tuple[torch.Tensor, torch.Tensor, int]:
     """Calculate the assistant token mask with `chat_template`.
 
@@ -35,6 +36,7 @@ def tokenize_and_mask_messages_hf(
         tools=tools,
         chat_template=chat_template,
         add_generation_prompt=False,
+        enable_thinking=enable_thinking,
         padding=False,
         truncation=True,
         return_tensors="pt",
@@ -52,6 +54,7 @@ def tokenize_and_mask_messages_default(
     messages: List[dict],
     tools: Optional[List[dict]] = None,
     chat_template: Optional[str] = None,
+    enable_thinking: Optional[bool] = None,
 ) -> Tuple[torch.Tensor, torch.Tensor, int]:
     """Calculate the assistant token mask.
 
@@ -78,6 +81,7 @@ def tokenize_and_mask_messages_default(
         tools=tools,
         chat_template=chat_template,
         add_generation_prompt=False,
+        enable_thinking=enable_thinking,
         padding=False,
         truncation=True,
         return_tensors="pt",
@@ -91,6 +95,7 @@ def tokenize_and_mask_messages_default(
                 tools=tools,
                 chat_template=chat_template,
                 add_generation_prompt=True,
+                enable_thinking=enable_thinking,
                 padding=False,
                 truncation=True,
                 return_tensors="pt",
@@ -102,6 +107,7 @@ def tokenize_and_mask_messages_default(
                 tools=tools,
                 chat_template=chat_template,
                 add_generation_prompt=False,
+                enable_thinking=enable_thinking,
                 padding=False,
                 truncation=True,
                 return_tensors="pt",

trinity/common/models/vllm_model.py

@@ -354,6 +354,7 @@ async def convert_messages_to_experience(
             messages=messages,
             tools=tools,
             chat_template=self.chat_template,
+            enable_thinking=self.enable_thinking,
         )  # (seq_length, ), (seq_length, )
         logprobs = await self.logprobs(token_ids=token_ids.tolist())  # (seq_length - 1,)
         return Experience(