Update app.py

app.py

@@ -10,11 +10,7 @@ OUTPUT_DIR = "train_output"
 ZIP_FILE = "python_ai_trained_model.zip"
 LOG_FILE = "train_log.txt"
 
-# Shared states
-start_time = None
-end_time = None
-
-# Function: Zip the trained model
+# Zip function
 def zip_trained_model():
     with zipfile.ZipFile(ZIP_FILE, 'w', zipfile.ZIP_DEFLATED) as zipf:
         for root, _, files in os.walk(OUTPUT_DIR):
@@ -23,66 +19,83 @@ def zip_trained_model():
                 arcname = os.path.relpath(filepath, OUTPUT_DIR)
                 zipf.write(filepath, arcname)
 
-# Function: Tail the logs from training
-def tail_logs():
+# Tail logs function
+def tail_logs(n=20):
     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)
+        return ''.join(f.readlines()[-n:])
 
-# Background training function
+# Background training runner
 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."
+    status_box.value = "🚀 Training started..."
     time_box.value = "Training in progress..."
     log_box.value = ""
 
-    # Run training and redirect stdout/stderr to log file
+    # Create log file
     with open(LOG_FILE, "w") as log:
-        process = subprocess.Popen(["python", "train.py"], stdout=log, stderr=log)
+        log.write("🚀 Launching train.py...\n")
+
+    # Start training
+    with open(LOG_FILE, "a") as log:
+        process = subprocess.Popen(["python", "train.py"], stdout=log, stderr=subprocess.STDOUT)
         while process.poll() is None:
-            log_box.value = tail_logs()
+            log_box.update(value=tail_logs())
             time.sleep(5)
 
-    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."
+    # Check exit status
+    if process.returncode != 0:
+        status_box.value = f"❌ Training failed with exit code {process.returncode}"
+        log_box.value = tail_logs()
+        return
 
-    status_box.value = "🔄 Compressing trained model..."
-    zip_trained_model()
+    # Training success
+    elapsed = round(time.time() - start_time, 2)
+    time_box.value = f"✅ Completed in {elapsed // 60:.0f} min {elapsed % 60:.0f} sec"
+    status_box.value = "🔄 Compressing model..."
 
+    # Zip it
+    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.value = ZIP_FILE
     download_file.visible = True
-    status_box.value = "✅ Done! You can now download your trained model."
+    status_box.value = "✅ Training complete. Download below."
 
-# Trigger button click
+# Button trigger
 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..."
+    return "Training process started."
 
 # Gradio UI
 with gr.Blocks() as demo:
     gr.Markdown("## 🧠 Python AI Trainer (StarCoder 7B)")
-    gr.Markdown("Train a custom Python AI that can write, fix, and explain Python code.")
+    gr.Markdown("Train your Python AI with 1 click. Watch logs. Download model when done.")
 
     with gr.Row():
         train_btn = gr.Button("🚀 Start Training")
         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)
+    log_box = gr.Textbox(label="Live Training Logs", lines=20, interactive=False, value="")
     download_file = gr.File(label="📥 Download Trained Model (.zip)", visible=False)
 
-    train_btn.click(fn=start_training, inputs=[status_box, time_box, download_file, model_size_box, log_box], outputs=[status_box])
+    train_btn.click(
+        fn=start_training,
+        inputs=[status_box, time_box, download_file, model_size_box, log_box],
+        outputs=[status_box]
+    )
+
+    status_box.render()
+    time_box.render()
+    model_size_box.render()
+    log_box.render()
+    download_file.render()
 
 demo.launch()