Fix Qwen3 thinking mode + increase max_new_tokens: training/train_sft.py

training/train_sft.py

@@ -96,7 +96,7 @@ class SftConfig:
     lr: float = 5e-5
     grad_clip: float = 1.0
     max_prompt_tokens: int = 2048
-    max_new_tokens: int = 96
+    max_new_tokens: int = 384
     bf16: bool = torch.cuda.is_available()
     use_4bit: bool = True
     gradient_checkpointing: bool = True
@@ -181,11 +181,15 @@ def build_sft_examples(
         for action_dict in _expert_actions(env, task_id):
             user_msg = build_observation_prompt(obs, is_first=(len(messages) == 1))
             messages.append({"role": "user", "content": user_msg})
-            prompt_text = tokenizer.apply_chat_template(
-                messages,
-                tokenize=False,
-                add_generation_prompt=True,
-            )
+            try:
+                prompt_text = tokenizer.apply_chat_template(
+                    messages, tokenize=False, add_generation_prompt=True,
+                    enable_thinking=False,
+                )
+            except TypeError:
+                prompt_text = tokenizer.apply_chat_template(
+                    messages, tokenize=False, add_generation_prompt=True,
+                )
             completion_text = json.dumps(action_dict, ensure_ascii=False)
             prompt_ids = tokenizer(
                 prompt_text,