project-llm-trainer 0.5.8__py3-none-any.whl → 0.5.9__py3-none-any.whl

llm_trainer/generate_utils.py

@@ -143,8 +143,7 @@ def _generate(
 
     with torch.inference_mode():
         for _ in range(max_new_tokens):
-            # 是否需要截取？？
-            t = tokens[:, -max_position_embeddings:]
+            t = tokens # tokens[:, -max_position_embeddings:]
             with ctx:
                 result = model(
                     t,
@@ -202,7 +201,7 @@ def _generate(
 def _streaming_generate(
         model: torch.nn.Module,
         *,
-        prompt: str,
+        prompt: Union[str, torch.Tensor],
         max_position_embeddings: int,
         max_new_tokens: int,
         temperature: Optional[float] = 1.0,
@@ -214,7 +213,11 @@ def _streaming_generate(
         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)
+
+    if isinstance(prompt, torch.Tensor):
+        encoded_tokens = prompt.to(device)
+    else:
+        encoded_tokens = TrainerTools().tokenizer.encode(prompt, unsqueeze=True, covert_tensor=True).to(device)
 
     generate_text_iterator = _generate(
         model=model,
@@ -237,7 +240,7 @@ def _streaming_generate(
 def streaming_generate(
         model: torch.nn.Module,
         *,
-        prompt: str,
+        prompt: Union[str, torch.Tensor],
         max_position_embeddings: int,
         max_new_tokens: int,
         temperature: Optional[float] = 1.0,
@@ -246,7 +249,8 @@ 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
+        device: Union[str, torch.device, int] = None,
+        return_token: bool = False
 ):
     text_iterator = _streaming_generate(
         model=model,
@@ -264,13 +268,16 @@ def streaming_generate(
 
     for (token, is_full_result) in text_iterator:
         if not is_full_result:
-            yield TrainerTools().tokenizer.decode(token.squeeze(0))
+            if return_token:
+                yield token.squeeze(0)
+            else:
+                yield TrainerTools().tokenizer.decode(token.squeeze(0))
 
 
 def generate(
         model: torch.nn.Module,
         *,
-        prompt: str,
+        prompt: Union[str, torch.Tensor],
         max_position_embeddings: int,
         max_new_tokens: int,
         temperature: Optional[float] = 1.0,
@@ -279,7 +286,8 @@ 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
+        device: Union[str, torch.device, int] = None,
+        return_token: bool = False
 ):
     text_iterator = _streaming_generate(
         model=model,
@@ -297,7 +305,12 @@ def generate(
 
     for (token, is_full_result) in text_iterator:
         if is_full_result:
-            return TrainerTools().tokenizer.decode(token.squeeze(0))
+            if return_token:
+                return token.squeeze(0)
+            else:
+                return TrainerTools().tokenizer.decode(token.squeeze(0))
+
+    return None
 
 
 def batch_generate(

llm_trainer/train_configs.py

@@ -164,7 +164,7 @@ class KDConfig:
 
 @dataclass(kw_only=True)
 class EvalConfig:
-    max_new_tokens: int = 512
+    max_new_tokens: Optional[int] = None
     temperature: float = 1.0
     top_p: float = 0.95
     top_k: Optional[float] = None

{project_llm_trainer-0.5.8.dist-info → project_llm_trainer-0.5.9.dist-info}/METADATA

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

{project_llm_trainer-0.5.8.dist-info → project_llm_trainer-0.5.9.dist-info}/RECORD

@@ -4,7 +4,7 @@ llm_trainer/dataset.py,sha256=4QlOo0SFB5816BUYegQjgobUqTUMQvdmZMM_OEAMSjE,4347
 llm_trainer/dpo_trainer.py,sha256=1A_4QP2_xqM_YeqdXy-0RaMvEL80gim-pgnPQyHww9U,12052
 llm_trainer/ds_checkpoint.py,sha256=D092fkS1Up4QmpV9YCpqbSzfX_caCAeX-UiOrhOE1I8,1947
 llm_trainer/eval.py,sha256=fjASCILU3fSPJxo9cP3rIXEEnkc5ZlUyHqXlZtUiHrw,888
-llm_trainer/generate_utils.py,sha256=CbJ3mfAD6DkQ0GUHcJQ1AK02m-ocwmd-BPXEpiwvNNQ,14933
+llm_trainer/generate_utils.py,sha256=wrZoG2g7CsOyG4sb3px9vURHQFV6_9j5kQmpFc5A8yg,15335
 llm_trainer/grpo_trainer.py,sha256=sCYjvksdm9f7TpN23KXuCmua_8VFTZEfVEcflL89P_I,16058
 llm_trainer/log.py,sha256=LxqTGRNZUGMTSQCePRpk-rYyxSnSIbT4kOdP8Fbzr0M,462
 llm_trainer/loss.py,sha256=NZCQeUXnLSj__mmDflE8g89KgE0emAJXIab0IERCLno,6023
@@ -17,17 +17,17 @@ llm_trainer/scheduler.py,sha256=LAI_0VxClsIQkix0bRoduRD4vPfVuIZDhZgTAT_KK8k,4901
 llm_trainer/sft_trainer.py,sha256=gxQA7T1o1QGUsHp2CX1Qb_fO5LppBJuNbc0H4ixCYUA,1783
 llm_trainer/tokenizer.py,sha256=SSpgXtb0e1NtQqRW0gCq09TTZi47umggy-Fh5EMHKJg,6708
 llm_trainer/tools.py,sha256=yF17lp6oOfLe2XJeKDQ1juZcbv-6vFamJSLwEeArduA,2975
-llm_trainer/train_configs.py,sha256=c6bgivkkWRYcPD3NzI5uRItAUhZiIBgKVMuMgVFRnFo,7336
+llm_trainer/train_configs.py,sha256=guV8xkG5TSGvYwFvsQV_mA8mDHLLVhL5L0xo_WMsMME,7347
 llm_trainer/trainer.py,sha256=U26dZc22nByfTZUzKeEiqqYVexBzgw0ep7N0Z2zIcWI,26141
 llm_trainer/utils.py,sha256=LWNhyQ0NDEZ9mZtk2Ryvh6EulvHIaUGIflugSpqmeFI,6791
-project_llm_trainer-0.5.8.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
-project_llm_trainer-0.5.8.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
-project_llm_trainer-0.5.8.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
-project_llm_trainer-0.5.8.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
-project_llm_trainer-0.5.8.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
-project_llm_trainer-0.5.8.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
-project_llm_trainer-0.5.8.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
-project_llm_trainer-0.5.8.dist-info/METADATA,sha256=54q4Nl2EWMYwShSqS8cLLqZ0iJVraAvJCGz8QPiVMiE,195
-project_llm_trainer-0.5.8.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-project_llm_trainer-0.5.8.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
-project_llm_trainer-0.5.8.dist-info/RECORD,,
+project_llm_trainer-0.5.9.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.5.9.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
+project_llm_trainer-0.5.9.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
+project_llm_trainer-0.5.9.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
+project_llm_trainer-0.5.9.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.5.9.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.5.9.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
+project_llm_trainer-0.5.9.dist-info/METADATA,sha256=YfFvnbVfUyNMCByDKCJ1rB4Mj0uxGQ2wquSe4QKaiF4,195
+project_llm_trainer-0.5.9.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.5.9.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.5.9.dist-info/RECORD,,