Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments, TrainerCallback
-from datasets import load_dataset, Dataset, DatasetDict
+from datasets import load_dataset, DatasetDict
 import os
 import time
 
@@ -123,7 +123,7 @@ class CustomCallback(TrainerCallback):
         global progress_info
         total_steps = state.max_steps
         current_step = state.global_step
-        elapsed_time = time.time() - state.log_history[0]["epoch_time"]
+        elapsed_time = time.time() - state.log_history[0].get("epoch_time", time.time())  # デフォルト値を追加
         time_per_step = elapsed_time / (current_step + 1)
         remaining_steps = total_steps - current_step
         time_remaining = time_per_step * remaining_steps
@@ -149,6 +149,7 @@ with gr.Blocks() as demo:
         progress.update(value=progress_info["progress"])
         time_remaining.value = f"{progress_info['time_remaining']}秒" if progress_info['time_remaining'] else "待機中"
 
-    train_button.click(fn=train_and_deploy, inputs=[token_input, repo_input, license_input], outputs=output).then(fn=update_ui)
+    train_button.click(fn=train_and_deploy, inputs=[token_input, repo_input, license_input], outputs=output)
+    train_button.click(fn=update_ui, inputs=[], outputs=[status, progress, time_remaining])
 
 demo.launch()