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

llm_trainer/checkpoint.py

@@ -81,12 +81,13 @@ def load_checkpoint(
     else:
         checkpoint_name = os.environ.get('CHECKPOINT_NAME', DEFAULT_CHECKPOINT_NAME)
 
-        state_dict = torch.load(checkpoint_name, weights_only=True, map_location=device)
-        raw_model = model.module if isinstance(model, DDP) else model
-        raw_model.load_state_dict(state_dict['model_state_dict'])
+        if os.path.exists(checkpoint_name):
+            state_dict = torch.load(checkpoint_name, weights_only=True, map_location=device)
+            raw_model = model.module if isinstance(model, DDP) else model
+            raw_model.load_state_dict(state_dict['model_state_dict'])
 
-        if optimizer:
-            optimizer.load_state_dict(state_dict['optim_state_dict'])
+            if optimizer:
+                optimizer.load_state_dict(state_dict['optim_state_dict'])
 
 
 def load_checkpoint_for_eval(

llm_trainer/eval.py

@@ -15,12 +15,17 @@ def submit_gen_task(
         tokens_per_image
 ):
     log_dir = get_log_dir()
+    tokens = TrainerTools().tokenizer.encode(prompt, unsqueeze=True, covert_tensor=True)
+
+    max_new_tokens = eval_config.max_new_tokens
+    if not max_new_tokens:
+        max_new_tokens = max_position_embeddings - tokens.shape[-1]
 
     gen_result = generate(
         eval_model,
-        prompt=prompt,
+        prompt=tokens,
         max_position_embeddings=max_position_embeddings,
-        max_new_tokens=eval_config.max_new_tokens,
+        max_new_tokens=max_new_tokens,
         temperature=eval_config.temperature,
         k=eval_config.top_k,
         p=eval_config.top_p,

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

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

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

@@ -1,9 +1,9 @@
 llm_trainer/__init__.py,sha256=HWgtTEVeQSnZmEyYQm2K6eFEG4X2QAoigMlB5Z2tcXE,260
-llm_trainer/checkpoint.py,sha256=RoRlIB-Qtvl3MyY3g0FbEBHLpFRLnEMZLpEncXOLToQ,4242
+llm_trainer/checkpoint.py,sha256=Wh5CwceIajTgJ9i_mH3I1R9N2nOLFqVFmlEMkTiGcD4,4306
 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/eval.py,sha256=ZyUfSo2Q8P-lrCdPEnGkoo5pGubd0AabREK5eMISRII,1109
 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
@@ -20,14 +20,14 @@ llm_trainer/tools.py,sha256=yF17lp6oOfLe2XJeKDQ1juZcbv-6vFamJSLwEeArduA,2975
 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.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,,
+project_llm_trainer-0.5.11.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.5.11.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
+project_llm_trainer-0.5.11.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
+project_llm_trainer-0.5.11.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
+project_llm_trainer-0.5.11.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.5.11.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.5.11.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
+project_llm_trainer-0.5.11.dist-info/METADATA,sha256=RAhT8VLTlO4Oyr9ocDxBvhulYed1JXWX_5GWWkiJ7go,196
+project_llm_trainer-0.5.11.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.5.11.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.5.11.dist-info/RECORD,,