Update README.md

README.md

@@ -37,10 +37,10 @@ tokenizer = AutoTokenizer.from_pretrained(
 )
 device = next(model.model.parameters()).device
 
-# if enable_think
-# formatted_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt = True, enable_thinking = True)
-# if disable_think
-formatted_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt = True, enable_thinking = False)
+# if open think mode, use the following code
+formatted_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt = True, enable_thinking = True)
+# if close think mode, use the following code
+# formatted_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt = True, enable_thinking = False)
 
 input_ids = tokenizer.encode(formatted_prompt, return_tensors='pt').to(device)
 outputs = model.generate(
@@ -48,10 +48,10 @@ outputs = model.generate(
     max_new_tokens=1000,
     do_sample=True
 )
-# if enable think
-# ans = [i.split("<|im_start|> assistant\n", 1)[1].strip() for i in tokenizer.batch_decode(outputs)]
-# if disable think
-ans = [i.split("<|im_start|> assistant\n<think>\n\n</think>", 1)[1].strip() for i in tokenizer.batch_decode(outputs)]
+# if open think mode, use the following code
+ans = [i.split("<|im_start|> assistant\n", 1)[1].strip() for i in tokenizer.batch_decode(outputs)]
+# if close think mode, use the following code
+# ans = [i.split("<|im_start|> assistant\n<think>\n\n</think>", 1)[1].strip() for i in tokenizer.batch_decode(outputs)]
 ```
 
 <p align="center">