project-llm-trainer 0.5.0__py3-none-any.whl → 0.5.1__py3-none-any.whl

llm_trainer/generate_utils.py

@@ -107,8 +107,7 @@ def _generate(
         pixel_values: Optional[torch.Tensor] = None,
         tokens_per_image: int = -1,
         suppress_tokens: Optional[List[int]] = None,
-        device: Union[str, torch.device, int],
-        reasoning_budget: Optional[int] = None
+        device: Union[str, torch.device, int]
 ):
     """
     :param model:
@@ -142,27 +141,6 @@ def _generate(
     kv_cache: Optional[KVCache] = None
     generate_tokens = tokens.clone()
 
-    reasoning_start = TrainerTools().tokenizer.reasoning_start
-    reasoning_end = TrainerTools().tokenizer.reasoning_end
-
-    # --- 状态初始化 ---
-    in_reasoning_block = False
-    reasoning_step_count = 0
-    # “冷静期”标志位。当强制结束思考后，在下一步抑制<reasoning>的生成。
-    suppress_reasoning_start_next = False
-
-    if reasoning_budget is not None:
-        prompt_tokens = tokens[0]
-        start_indices = (prompt_tokens == reasoning_start).nonzero(as_tuple=True)[0]
-        end_indices = (prompt_tokens == reasoning_end).nonzero(as_tuple=True)[0]
-
-        last_start_idx = start_indices[-1].item() if len(start_indices) > 0 else -1
-        last_end_idx = end_indices[-1].item() if len(end_indices) > 0 else -1
-
-        if last_start_idx > last_end_idx:
-            in_reasoning_block = True
-            reasoning_step_count = len(prompt_tokens) - 1 - last_start_idx
-
     with torch.inference_mode():
         for _ in range(max_new_tokens):
             # 是否需要截取？？
@@ -182,23 +160,6 @@ def _generate(
             # (batch, vocab_size)
             logits = logits[:, -1, :]
 
-            # --- 推理预算逻辑 ---
-            force_end_reasoning_token = False
-            if reasoning_budget is not None:
-                # 检查是否需要在此步抑制 <reasoning>
-                should_suppress_this_step = suppress_reasoning_start_next
-                suppress_reasoning_start_next = False  # 立即重置标志位
-
-                # 修改: 检查是否超出预算
-                if in_reasoning_block and reasoning_step_count >= reasoning_budget:
-                    force_end_reasoning_token = True
-                    # 设置标志位，在下一步抑制 <reasoning>
-                    suppress_reasoning_start_next = True
-
-                # 如果上一轮设置了抑制标志，则在此轮执行抑制
-                if should_suppress_this_step:
-                    logits[:, reasoning_start] = -float("inf")
-
             # 抑制特殊token输出
             if suppress_tokens and len(suppress_tokens) != 0:
                 logits = _suppress_warper(logits, suppress_tokens)
@@ -214,10 +175,6 @@ def _generate(
             if p and 0 < p <= 1:
                 logits = _top_p_warper(logits, p)
 
-            if force_end_reasoning_token:
-                logits[:] = -float("inf")
-                logits[:, reasoning_end] = 0.0
-
             if multinomial:
                 prob = logits.softmax(dim=-1)
                 # 返回下标
@@ -226,18 +183,6 @@ def _generate(
                 # 返回下标
                 next_token = logits.argmax(dim=-1, keepdim=True)
 
-            if reasoning_budget is not None:
-                current_token_id = next_token.item()
-                if not in_reasoning_block and current_token_id == reasoning_start:
-                    in_reasoning_block = True
-                    reasoning_step_count = 0
-                elif in_reasoning_block:
-                    if current_token_id == reasoning_end:
-                        in_reasoning_block = False
-                        reasoning_step_count = 0
-                    else:
-                        reasoning_step_count += 1
-
             # token, is_full_result
             yield next_token, False
 
@@ -266,8 +211,7 @@ def _streaming_generate(
         pixel_values: Optional[torch.Tensor] = None,
         tokens_per_image: int = -1,
         suppress_tokens: Optional[List[int]] = None,
-        device: Union[str, torch.device, int] = None,
-        reasoning_budget: Optional[int] = None
+        device: Union[str, torch.device, int] = None
 ):
     device = TrainerTools().parallel.device if not device else device
     encoded_tokens = TrainerTools().tokenizer.encode(prompt, unsqueeze=True, covert_tensor=True).to(device)
@@ -283,8 +227,7 @@ def _streaming_generate(
         pixel_values=pixel_values,
         tokens_per_image=tokens_per_image,
         suppress_tokens=suppress_tokens,
-        device=device,
-        reasoning_budget=reasoning_budget
+        device=device
     )
 
     for (token, is_full_result) in generate_text_iterator:
@@ -303,8 +246,7 @@ def streaming_generate(
         pixel_values: Optional[torch.Tensor] = None,
         tokens_per_image: int = -1,
         suppress_tokens: Optional[List[int]] = None,
-        device: Union[str, torch.device, int] = None,
-        reasoning_budget: Optional[int] = None
+        device: Union[str, torch.device, int] = None
 ):
     text_iterator = _streaming_generate(
         model=model,
@@ -317,8 +259,7 @@ def streaming_generate(
         pixel_values=pixel_values,
         tokens_per_image=tokens_per_image,
         suppress_tokens=suppress_tokens,
-        device=device,
-        reasoning_budget=reasoning_budget
+        device=device
     )
 
     for (token, is_full_result) in text_iterator:
@@ -338,8 +279,7 @@ def generate(
         pixel_values: Optional[torch.Tensor] = None,
         tokens_per_image: int = -1,
         suppress_tokens: Optional[List[int]] = None,
-        device: Union[str, torch.device, int] = None,
-        reasoning_budget: Optional[int] = None
+        device: Union[str, torch.device, int] = None
 ):
     text_iterator = _streaming_generate(
         model=model,
@@ -352,8 +292,7 @@ def generate(
         suppress_tokens=suppress_tokens,
         pixel_values=pixel_values,
         tokens_per_image=tokens_per_image,
-        device=device,
-        reasoning_budget=reasoning_budget
+        device=device
     )
 
     for (token, is_full_result) in text_iterator:

llm_trainer/tokenizer.py

@@ -26,8 +26,8 @@ class Tokenizer:
         self.text_user = '<user>'
         self.text_assistant = '<assistant>'
 
-        self.text_reasoning_start = '<reasoning>'
-        self.text_reasoning_end = '</reasoning>'
+        self.text_think_start = '<think>'
+        self.text_think_end = '</think>'
 
         self.text_answer_start = '<answer>'
         self.text_answer_end = '</answer>'
@@ -47,8 +47,8 @@ class Tokenizer:
             additional_special_tokens = [
                 AddedToken(self.text_user, lstrip=False, rstrip=False),
                 AddedToken(self.text_assistant, lstrip=False, rstrip=False),
-                AddedToken(self.text_reasoning_start, lstrip=False, rstrip=False),
-                AddedToken(self.text_reasoning_end, lstrip=False, rstrip=False),
+                AddedToken(self.text_think_start, lstrip=False, rstrip=False),
+                AddedToken(self.text_think_end, lstrip=False, rstrip=False),
                 AddedToken(self.text_answer_start, lstrip=False, rstrip=False),
                 AddedToken(self.text_answer_end, lstrip=False, rstrip=False),
                 AddedToken(self.text_system, lstrip=False, rstrip=False),
@@ -69,8 +69,8 @@ class Tokenizer:
         self.user = self.tokenizer.convert_tokens_to_ids(self.text_user)
         self.assistant = self.tokenizer.convert_tokens_to_ids(self.text_assistant)
 
-        self.reasoning_start = self.tokenizer.convert_tokens_to_ids(self.text_reasoning_start)
-        self.reasoning_end = self.tokenizer.convert_tokens_to_ids(self.text_reasoning_end)
+        self.think_start = self.tokenizer.convert_tokens_to_ids(self.text_think_start)
+        self.think_end = self.tokenizer.convert_tokens_to_ids(self.text_think_end)
 
         self.answer_start = self.tokenizer.convert_tokens_to_ids(self.text_answer_start)
         self.answer_end = self.tokenizer.convert_tokens_to_ids(self.text_answer_end)
@@ -140,9 +140,9 @@ class Tokenizer:
                 {"role":"user", "content":"hello?"},
                 {"role":"assistant", "content":"hello"},
                 {"role":"user", "content":"hello hello?"},
-                {"role":"assistant", "reasoning":"thinking", "content":"hello hello"},
+                {"role":"assistant", "think":"thinking", "content":"hello hello"},
             ]
-            <system>{system_prompt}</s><user>hello?</s><assistant>hello</s><user>hello hello?</s><assistant><reasoning>thinking</reasoning><answer>hello hello</answer></s>
+            <system>{system_prompt}</s><user>hello?</s><assistant>hello</s><user>hello hello?</s><assistant><think>thinking</think><answer>hello hello</answer></s>
         """
 
         chat_template = ''
@@ -154,8 +154,8 @@ class Tokenizer:
                 if add_answer_tag_for_assistant and role == 'assistant':
                     content = f"{self.text_answer_start}{content}{self.text_answer_end}"
 
-                if 'reasoning' in conversation:
-                    content = f"{self.text_reasoning_start}{conversation['reasoning']}{self.text_reasoning_end}{content}"
+                if 'think' in conversation:
+                    content = f"{self.text_think_start}{conversation['think']}{self.text_think_end}{content}"
 
                 chat_template = f"{chat_template}{support_roles[role]}{content}{self.text_end}"
 

{project_llm_trainer-0.5.0.dist-info → project_llm_trainer-0.5.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: project_llm_trainer
-Version: 0.5.0
+Version: 0.5.1
 Summary: LLM and VLM trainer
 Author: qibin
 Author-email: qibin0506@gmail.com

{project_llm_trainer-0.5.0.dist-info → project_llm_trainer-0.5.1.dist-info}/RECORD

@@ -4,7 +4,7 @@ llm_trainer/dataset.py,sha256=4QlOo0SFB5816BUYegQjgobUqTUMQvdmZMM_OEAMSjE,4347
 llm_trainer/dpo_trainer.py,sha256=wMREatLt0I8Ajdm_sI2U8Zj-IN1L6txP9s_tH1oI3-s,11431
 llm_trainer/ds_checkpoint.py,sha256=wz48HoLBBt8QGO1tXfvJwrXoiGtPG_gjwHfEqARllso,2175
 llm_trainer/eval.py,sha256=fjASCILU3fSPJxo9cP3rIXEEnkc5ZlUyHqXlZtUiHrw,888
-llm_trainer/generate_utils.py,sha256=2MoEGEpoTzx7khO3dPcC2akFLyjtbFFpdJtuB_QQ3OY,17708
+llm_trainer/generate_utils.py,sha256=CbJ3mfAD6DkQ0GUHcJQ1AK02m-ocwmd-BPXEpiwvNNQ,14933
 llm_trainer/grpo_trainer.py,sha256=qiC3KwxYPSB9UKqyk4eSRvORP3b6GM-2ozqI8u3QvI0,15568
 llm_trainer/log.py,sha256=LxqTGRNZUGMTSQCePRpk-rYyxSnSIbT4kOdP8Fbzr0M,462
 llm_trainer/loss.py,sha256=NZCQeUXnLSj__mmDflE8g89KgE0emAJXIab0IERCLno,6023
@@ -15,19 +15,19 @@ llm_trainer/parallel_none.py,sha256=TG6Pm829Dg-yQu-97O-EHV3FCARBlNcP47KkGFAs16E,
 llm_trainer/partition_utils.py,sha256=xzv8kwlbKp3dai2pBwX89gN5ymeHk1bGbTkGru5H-UM,5167
 llm_trainer/scheduler.py,sha256=Xz8HhwoRMjRe41sf_NHhpZfkTlEs0I2MYusvMY6hCVw,3531
 llm_trainer/sft_trainer.py,sha256=gxQA7T1o1QGUsHp2CX1Qb_fO5LppBJuNbc0H4ixCYUA,1783
-llm_trainer/tokenizer.py,sha256=A7TYYUbtPf75kjCvWP7yBui4xZBObMk2aPem62YpwpY,6776
+llm_trainer/tokenizer.py,sha256=SSpgXtb0e1NtQqRW0gCq09TTZi47umggy-Fh5EMHKJg,6708
 llm_trainer/tools.py,sha256=yF17lp6oOfLe2XJeKDQ1juZcbv-6vFamJSLwEeArduA,2975
 llm_trainer/train_configs.py,sha256=m57W71SI5VCCU9aJ_nJkB-3AJrSGiNXmV28rdpuYmLg,7332
 llm_trainer/trainer.py,sha256=zTJVyY1cAjJdTkyXCOy2ZPVP18SOMLdWhD54Mz2JRe4,25314
 llm_trainer/utils.py,sha256=LWNhyQ0NDEZ9mZtk2Ryvh6EulvHIaUGIflugSpqmeFI,6791
-project_llm_trainer-0.5.0.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
-project_llm_trainer-0.5.0.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
-project_llm_trainer-0.5.0.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
-project_llm_trainer-0.5.0.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
-project_llm_trainer-0.5.0.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
-project_llm_trainer-0.5.0.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
-project_llm_trainer-0.5.0.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
-project_llm_trainer-0.5.0.dist-info/METADATA,sha256=YDj-N4VL8O_AqNanwfU6Yt38J97p3RgtUSzmwl0Y-GM,195
-project_llm_trainer-0.5.0.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-project_llm_trainer-0.5.0.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
-project_llm_trainer-0.5.0.dist-info/RECORD,,
+project_llm_trainer-0.5.1.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.5.1.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
+project_llm_trainer-0.5.1.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
+project_llm_trainer-0.5.1.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
+project_llm_trainer-0.5.1.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.5.1.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.5.1.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
+project_llm_trainer-0.5.1.dist-info/METADATA,sha256=x-Bobn0EH7wyKznJydUeVLK9sdIrkBmDYDbEpyG4pKc,195
+project_llm_trainer-0.5.1.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.5.1.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.5.1.dist-info/RECORD,,