Update app.py

app.py

@@ -2,8 +2,27 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
 from datasets import load_dataset, Dataset, DatasetDict
 import os
+import time
+
+# トレーニングの進行状況を格納するグローバル変数
+progress_info = {
+    "status": "待機中",
+    "progress": 0,
+    "time_remaining": None
+}
+
+def update_progress(trainer, epoch, step, total_steps, time_remaining):
+    global progress_info
+    progress_info["status"] = f"エポック {epoch + 1} / {trainer.args.num_train_epochs}, ステップ {step + 1} / {total_steps}"
+    progress_info["progress"] = (step + 1) / total_steps
+    progress_info["time_remaining"] = time_remaining
 
 def train_and_deploy(write_token, repo_name, license_text):
+    global progress_info
+    progress_info["status"] = "トレーニング開始"
+    progress_info["progress"] = 0
+    progress_info["time_remaining"] = None
+
     # トークンを環境変数に設定
     os.environ['HF_WRITE_TOKEN'] = write_token
     
@@ -62,16 +81,50 @@ def train_and_deploy(write_token, repo_name, license_text):
         args=training_args,
         train_dataset=tokenized_datasets["train"],
         eval_dataset=tokenized_datasets["test"],
+        callbacks=[CustomCallback()]
     )
     
     # トレーニング実行
+    start_time = time.time()
     trainer.train()
+    end_time = time.time()
+    total_time = end_time - start_time
+    progress_info["status"] = f"トレーニング完了（所要時間: {total_time:.2f}秒）"
+    progress_info["progress"] = 1
+    progress_info["time_remaining"] = 0
     
     # モデルをHugging Face Hubにプッシュ
     trainer.push_to_hub()
-
+    
     return f"モデルが'{repo_name}'リポジトリにデプロイされました！"
 
+class CustomCallback(TrainerCallback):
+    def on_train_begin(self, args, state, control, **kwargs):
+        global progress_info
+        progress_info["status"] = "トレーニング開始"
+        progress_info["progress"] = 0
+        progress_info["time_remaining"] = None
+
+    def on_step_begin(self, args, state, control, **kwargs):
+        global progress_info
+        total_steps = state.num_train_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
+        progress_info["time_remaining"] = None
+
+    def on_step_end(self, args, state, control, **kwargs):
+        global progress_info
+        total_steps = state.num_train_steps
+        current_step = state.global_step
+        elapsed_time = time.time() - state.log_history[0]["epoch_time"]
+        time_per_step = elapsed_time / (current_step + 1)
+        remaining_steps = total_steps - current_step
+        time_remaining = time_per_step * remaining_steps
+        progress_info["status"] = f"エポック {state.epoch + 1} / {args.num_train_epochs}, ステップ {current_step + 1} / {total_steps}"
+        progress_info["progress"] = (current_step + 1) / total_steps
+        progress_info["time_remaining"] = f"{time_remaining:.2f}秒"
+
 # Gradio UI
 with gr.Blocks() as demo:
     gr.Markdown("### pythia トレーニングとデプロイ")
@@ -79,8 +132,17 @@ with gr.Blocks() as demo:
     repo_input = gr.Textbox(label="リポジトリ名", placeholder="デプロイするリポジトリ名を入力してください...")
     license_input = gr.Textbox(label="ライセンス", placeholder="ライセンス情報を入力してください...")
     output = gr.Textbox(label="出力")
+    progress = gr.Progress(track_tqdm=True)
+    status = gr.Textbox(label="ステータス", value="待機中")
+    time_remaining = gr.Textbox(label="残り時間", value="待機中")
     train_button = gr.Button("デプロイ")
 
-    train_button.click(fn=train_and_deploy, inputs=[token_input, repo_input, license_input], outputs=output)
+    def update_ui():
+        global progress_info
+        status.update(value=progress_info["status"])
+        progress.update(value=progress_info["progress"])
+        time_remaining.update(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)
 
 demo.launch()