Update app.py: auto-start training, file-based logging

app.py

@@ -2,39 +2,41 @@
 """
 Gradio App for EeshaAI/Zeeb Training Space
 ==========================================
-Provides a web UI to trigger LoRA fine-tuning of OLMo 2 1B Instruct
-on tokenized video data, then push the trained model to EeshaAI/zeeb.
+Auto-starts LoRA fine-tuning on Space boot.
+The UI shows real-time training progress from the log file.
 """
 
 import os
-import sys
-import io
+import time
+import threading
 import gradio as gr
 
+LOG_FILE = "/tmp/training_log.txt"
 
-def run_training():
-    """Run the training pipeline, capturing all output."""
-    # Capture all prints and logs
-    old_stdout = sys.stdout
-    old_stderr = sys.stderr
-    sys.stdout = buffer = io.StringIO()
-    sys.stderr = buffer
 
-    log_output = ""
+def start_training_background():
+    """Start training in a background thread on Space startup."""
+    from train_on_hf_spaces import run_training_to_file
+    run_training_to_file(LOG_FILE)
 
+
+def get_log():
+    """Read the current training log."""
     try:
-        from train_on_hf_spaces import train
-        for log_msg in train("tokenized_dataset.json"):
-            log_output += log_msg
-    except Exception as e:
-        import traceback
-        log_output += f"\n❌ ERROR: {e}\n"
-        log_output += traceback.format_exc()
-    finally:
-        sys.stdout = old_stdout
-        sys.stderr = old_stderr
+        with open(LOG_FILE, "r") as f:
+            return f.read()
+    except FileNotFoundError:
+        return "⏳ Training has not started yet. Please wait..."
+
+
+def refresh_log():
+    """Refresh button callback."""
+    return get_log()
 
-    return log_output
+
+# Auto-start training on Space boot
+training_thread = threading.Thread(target=start_training_background, daemon=True)
+training_thread.start()
 
 
 with gr.Blocks(
@@ -45,36 +47,30 @@ with gr.Blocks(
     gr.Markdown(
         """
         # 🎬 Zeeb — Video-LLM Trainer
-        Fine-tune **OLMo 2 1B Instruct** with **LoRA (r=4)** to generate video tokens.
-        Trained model is automatically pushed to [EeshaAI/zeeb](https://huggingface.co/EeshaAI/zeeb).
+        Fine-tuning **OLMo 2 1B Instruct** with **LoRA (r=4)** to generate video tokens.
+        Trained model will be pushed to [EeshaAI/zeeb](https://huggingface.co/EeshaAI/zeeb).
+
+        Training **starts automatically** when this Space boots.
         """
     )
 
-    train_btn = gr.Button("🚀 Start Training", variant="primary", size="lg")
+    with gr.Row():
+        refresh_btn = gr.Button("🔄 Refresh Log", variant="primary")
+        auto_refresh = gr.Checkbox(label="Auto-refresh (every 10s)", value=True)
 
     logbox = gr.Textbox(
         label="Training Log",
+        value=lambda: get_log(),
         lines=30,
         max_lines=200,
         interactive=False,
         show_copy_button=True,
     )
 
-    gr.Markdown(
-        """
-        ### What happens when you click "Start Training"?
-        1. 📦 Downloads **OLMo 2 1B Instruct** from HuggingFace
-        2. 🔤 Expands vocabulary with **1,024 visual tokens** (`<v_0>` ... `<v_1023>`)
-        3. 🔧 Applies **LoRA r=4** to q_proj & v_proj (minimal memory)
-        4. 🔥 Trains for **3 epochs** on the tokenized video dataset
-        5. 🔀 Merges LoRA weights back into the base model
-        6. 🚀 Pushes the merged model to **EeshaAI/zeeb**
-
-        ⚠️ Training on CPU takes time (~10-30 min depending on dataset size).
-        """
-    )
+    refresh_btn.click(fn=refresh_log, outputs=logbox)
 
-    train_btn.click(fn=run_training, outputs=logbox)
+    # Auto-refresh every 10 seconds
+    demo.load(fn=refresh_log, outputs=logbox, every=10)
 
 
 if __name__ == "__main__":