Adding the enable_thinking args

Author: LakoMoor

README.md

@@ -215,6 +215,7 @@ PYTHONPATH="${PYTHONPATH}:src/" poetry run accelerate launch --config_file accel
     - "lm_head"
   lora_r: 256
   lora_alpha: 256
+  enable_thinking: False # set True if the model and dataset supports thinking
   assistant_message_template: "<|start_header_id|>assistant<|end_header_id|>\n\n"
   custom_pad_token: "<|reserved_special_token_0|>"
   custom_eos_token: "<|eot_id|>"

scripts/model_training/sft.py

@@ -94,7 +94,8 @@ def process_row(row, add_gen_prompt=False):
         constructed_prompt = tokenizer.apply_chat_template(
             system_message + history,
             tokenize=False,
-            add_generation_prompt=add_gen_prompt
+            add_generation_prompt=add_gen_prompt,
+            enable_thinking=args.enable_thinking
         )
         if tokenizer.bos_token is not None:
             if constructed_prompt.startswith(tokenizer.bos_token):  # Remove extra bos token

scripts/prompts_training/sft.py

@@ -100,7 +100,8 @@ def process_row(row, add_gen_prompt=False):
         constructed_prompt = tokenizer.apply_chat_template(
             system_message + history,
             tokenize=False,
-            add_generation_prompt=add_gen_prompt
+            add_generation_prompt=add_gen_prompt,
+            enable_thinking=args.enable_thinking
         )
         if tokenizer.bos_token is not None:
             if constructed_prompt.startswith(tokenizer.bos_token):  # Remove extra bos token

src/configs/additional/sft_args.py

@@ -17,6 +17,9 @@ class SFTScriptArguments(CommonScriptArguments):
             "help": "Will use system prompt if there is no one in dialogue, set to None to disable"
         },
     )
+    enable_thinking: bool | None = field(
+        default=False, metadata={"help": "Do train with thinking"}
+    )
     train_only_on_completions: bool | None = field(
         default=True, metadata={"help": "Do train only on completions or not"}
     )