Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
+from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments, TrainerCallback
 from datasets import load_dataset, Dataset, DatasetDict
 import os
 import time
@@ -107,7 +107,7 @@ class CustomCallback(TrainerCallback):
 
     def on_step_begin(self, args, state, control, **kwargs):
         global progress_info
-        total_steps = state.num_train_steps
+        total_steps = state.max_steps
         current_step = state.global_step
         progress_info["status"] = f"エポック {state.epoch + 1} / {args.num_train_epochs}, ステップ {current_step + 1} / {total_steps}"
         progress_info["progress"] = (current_step + 1) / total_steps
@@ -115,7 +115,7 @@ class CustomCallback(TrainerCallback):
 
     def on_step_end(self, args, state, control, **kwargs):
         global progress_info
-        total_steps = state.num_train_steps
+        total_steps = state.max_steps
         current_step = state.global_step
         elapsed_time = time.time() - state.log_history[0]["epoch_time"]
         time_per_step = elapsed_time / (current_step + 1)