Update app.py: better error handling, show_error=True

app.py

@@ -6,24 +6,40 @@ 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.
 """
 
+import os
+import sys
+import io
 import gradio as gr
-from train_on_hf_spaces import train
 
 
 def run_training():
-    """Run the training pipeline and stream logs to the UI."""
-    for log_msg in train("tokenized_dataset.json"):
-        yield log_msg
+    """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 = ""
+
+    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
+
+    return log_output
 
 
 with gr.Blocks(
     title="Zeeb — Video-LLM Trainer",
     theme=gr.themes.Soft(),
-    css="""
-        #logbox { font-family: 'Courier New', monospace; font-size: 13px; line-height: 1.5; }
-        #title { text-align: center; }
-        .contain { max-width: 800px; margin: auto; }
-    """
 ) as demo:
 
     gr.Markdown(
@@ -31,18 +47,15 @@ with gr.Blocks(
         # 🎬 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).
-        """,
-        elem_id="title",
+        """
     )
 
-    with gr.Row():
-        train_btn = gr.Button("🚀 Start Training", variant="primary", size="lg")
+    train_btn = gr.Button("🚀 Start Training", variant="primary", size="lg")
 
     logbox = gr.Textbox(
         label="Training Log",
-        elem_id="logbox",
-        lines=25,
-        max_lines=100,
+        lines=30,
+        max_lines=200,
         interactive=False,
         show_copy_button=True,
     )
@@ -65,4 +78,4 @@ with gr.Blocks(
 
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)