fix transforrmers api change at 5.2.0

Author: UbeCc

slime/utils/mask_utils.py

@@ -36,11 +36,15 @@ def get_system_message_length(self) -> tuple[int, int]:
         chat_template_token_ids = self.tokenizer(chat_template_token, add_special_tokens=False)["input_ids"]
         idx_1, idx_2 = self.find_all_sublist_indices(chat_template_token_ids, raw_token_ids)
         end_interval = len(chat_template_token_ids) - len(raw_token_ids) - idx_2
-        gen_token_length = len(
-            self.tokenizer.apply_chat_template(
-                test_messages, add_special_tokens=False, tokenize=True, add_generation_prompt=True
-            )
-        ) - len(chat_template_token_ids)
+        
+        gen_prompt_token_ids = self.tokenizer.apply_chat_template(
+            test_messages, add_special_tokens=False, tokenize=True, add_generation_prompt=True
+        )
+        # Handle transformers 5.2.0+ API change: apply_chat_template now returns dict when tokenize=True
+        if not isinstance(gen_prompt_token_ids, list):
+            gen_prompt_token_ids = gen_prompt_token_ids["input_ids"]
+        
+        gen_token_length = len(gen_prompt_token_ids) - len(chat_template_token_ids)
 
         system_message_length = idx_1 - ((idx_2 - idx_1) - end_interval - len(raw_token_ids))
         return system_message_length, gen_token_length
@@ -57,6 +61,10 @@ def gen_multi_turn_loss_mask_qwen(
             else:
                 message_ids = self.tokenizer.apply_chat_template([message], tokenize=True)
 
+            # Handle transformers 5.2.0+ API change: apply_chat_template now returns dict when tokenize=True
+            if not isinstance(message_ids, list):
+                message_ids = message_ids["input_ids"]
+            
             if message["role"] != "system" and i > 0:
                 message_ids = message_ids[self.system_message_length :]
 
@@ -81,15 +89,25 @@ def gen_multi_turn_loss_mask_qwen3(
 
         prefix_message = {"role": "user", "content": "FOR CALCULATING LOSS MASK ONLY"}
         prefix_token_ids = self.tokenizer.apply_chat_template([prefix_message], tokenize=True)
+        
+        # Handle transformers 5.2.0+ API change: apply_chat_template now returns dict when tokenize=True
+        if not isinstance(prefix_token_ids, list):
+            prefix_token_ids = prefix_token_ids["input_ids"]
 
         for i, message in enumerate(messages):
             if i == 0:
                 tailed_message_ids = self.tokenizer.apply_chat_template(
                     [message, prefix_message], tokenize=True, tools=tools
                 )
+                # Handle transformers 5.2.0+ API change
+                if not isinstance(tailed_message_ids, list):
+                    tailed_message_ids = tailed_message_ids["input_ids"]
                 message_ids = tailed_message_ids[: -len(prefix_token_ids)]
             else:
                 prefixed_message_ids = self.tokenizer.apply_chat_template([prefix_message, message], tokenize=True)
+                # Handle transformers 5.2.0+ API change
+                if not isinstance(prefixed_message_ids, list):
+                    prefixed_message_ids = prefixed_message_ids["input_ids"]
                 message_ids = prefixed_message_ids[len(prefix_token_ids) :]
 
             if message["role"] != "system" and i > 0:
@@ -181,3 +199,12 @@ def get_text_from_loss_mask(self, token_ids: list[int], loss_masks: list[int]) -
             selected_texts.append(self.tokenizer.decode(current_tokens))
 
         return selected_texts
+
+if __name__ == "__main__":
+    tokenizer = AutoTokenizer.from_pretrained("/workspace/Qwen3.5-35B-A3B")
+    mask_utils = MultiTurnLossMaskGenerator(tokenizer, tokenizer_type="qwen")
+    messages = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}]
+    tools = [{"type": "function", "function": {"name": "get_weather", "description": "Get the weather", "parameters": {"type": "object", "properties": {"city": {"type": "string", "description": "The city to get the weather for"}}}}}]
+    token_ids, loss_mask = mask_utils.get_loss_mask(messages, tools=tools)
+    for i in range(len(token_ids)):
+        print(f'Token: {repr(tokenizer.decode(token_ids[i]))}, Loss Mask: {loss_mask[i]}')