Update app.py

app.py

@@ -1,49 +1,87 @@
 import os
-import subprocess
-import gradio as gr
-import threading
 import time
 import zipfile
+import gradio as gr
+import subprocess
+import threading
 
-output_path = "train_output"
-zipped_model = "python_ai_trained_model.zip"
-status = gr.Textbox(label="Status", value="Ready", interactive=False)
-download_link = gr.File(label="Download Trained Model", visible=False)
-
-def run_training():
-    global zipped_model
-    status.value = "Training started... this may take a while (15–60+ mins)."
+# Paths
+OUTPUT_DIR = "train_output"
+ZIP_FILE = "python_ai_trained_model.zip"
+LOG_FILE = "train_log.txt"
 
-    # Run train.py
-    subprocess.run(["python", "train.py"], check=True)
+# Shared states
+start_time = None
+end_time = None
 
-    # Compress trained model
-    status.value = "Training complete. Compressing model..."
-    with zipfile.ZipFile(zipped_model, 'w', zipfile.ZIP_DEFLATED) as zipf:
-        for root, _, files in os.walk(output_path):
+# Function: Zip the trained model
+def zip_trained_model():
+    with zipfile.ZipFile(ZIP_FILE, 'w', zipfile.ZIP_DEFLATED) as zipf:
+        for root, _, files in os.walk(OUTPUT_DIR):
             for file in files:
                 filepath = os.path.join(root, file)
-                arcname = os.path.relpath(filepath, output_path)
+                arcname = os.path.relpath(filepath, OUTPUT_DIR)
                 zipf.write(filepath, arcname)
 
-    # Move zip to Gradio-visible path
-    status.value = "Done. Model ready for download."
-    download_link.visible = True
+# Function: Tail the logs from training
+def tail_logs():
+    if not os.path.exists(LOG_FILE):
+        return ""
+    with open(LOG_FILE, 'r') as f:
+        lines = f.readlines()[-20:]  # last 20 lines
+        return ''.join(lines)
+
+# Background training function
+def run_training(status_box, time_box, download_file, model_size_box, log_box):
+    global start_time, end_time
+
+    start_time = time.time()
+    status_box.value = "🚀 Training started... Please wait. This can take 30–90+ minutes."
+    time_box.value = "Training in progress..."
+    log_box.value = ""
+
+    # Run training and redirect stdout/stderr to log file
+    with open(LOG_FILE, "w") as log:
+        process = subprocess.Popen(["python", "train.py"], stdout=log, stderr=log)
+        while process.poll() is None:
+            log_box.value = tail_logs()
+            time.sleep(5)
 
-def start_training():
-    training_thread = threading.Thread(target=run_training)
-    training_thread.start()
-    return "Training in progress...", None
+    end_time = time.time()
+    elapsed = round(end_time - start_time, 2)
+    time_box.value = f"✅ Training completed in {elapsed // 60:.0f} min {elapsed % 60:.0f} sec."
 
+    status_box.value = "🔄 Compressing trained model..."
+    zip_trained_model()
+
+    size_mb = round(os.path.getsize(ZIP_FILE) / (1024 * 1024), 2)
+    model_size_box.value = f"📦 Model Size: {size_mb} MB"
+
+    download_file.visible = True
+    status_box.value = "✅ Done! You can now download your trained model."
+
+# Trigger button click
+def start_training(status_box, time_box, download_file, model_size_box, log_box):
+    thread = threading.Thread(target=run_training, args=(status_box, time_box, download_file, model_size_box, log_box))
+    thread.start()
+    return "Training process has started..."
+
+# Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## 🧠 Python AI Model Trainer (StarCoder 7B)")
-    gr.Markdown("Click the button below to start fine-tuning your custom Python AI. After training, download the model and use it anywhere.")
-    
+    gr.Markdown("## 🧠 Python AI Trainer (StarCoder 7B)")
+    gr.Markdown("Train a custom Python AI that can write, fix, and explain Python code.")
+
     with gr.Row():
         train_btn = gr.Button("🚀 Start Training")
-        train_btn.click(start_training, outputs=[status, download_link])
+        status_box = gr.Textbox(label="Status", value="Ready", interactive=False)
+    
+    with gr.Row():
+        time_box = gr.Textbox(label="Training Time", interactive=False)
+        model_size_box = gr.Textbox(label="Final Model Size", interactive=False)
+
+    log_box = gr.Textbox(label="Live Training Logs", lines=20, interactive=False)
+    download_file = gr.File(label="📥 Download Trained Model (.zip)", value=ZIP_FILE, visible=False)
 
-    status.render()
-    download_link.render()
+    train_btn.click(fn=start_training, inputs=[status_box, time_box, download_file, model_size_box, log_box], outputs=[status_box])
 
 demo.launch()