Update app.py

app.py

@@ -1,39 +1,73 @@
 import gradio as gr
-import subprocess
 import os
+import subprocess
+from pathlib import Path
+
+def check_training_status():
+    results_dir = Path("/app/results")
+    if not results_dir.exists():
+        return "Training hasn't started yet."
+    
+    iterations = len(list(results_dir.glob("iter_*")))
+    return f"Completed {iterations} training iterations."
 
-def run_self_lengthen(model_path, instruct_count, max_iter):
+def start_training(model_path, instruct_count, max_iter):
+    os.environ["MODEL_PATH"] = model_path
+    os.environ["INSTRUCT_COUNT"] = str(instruct_count)
+    os.environ["MAX_ITER"] = str(max_iter)
+    
     try:
-        process = subprocess.Popen([
-            "bash", 
-            "/app/qwen/run.sh",
-            f"--base_model={model_path}",
-            f"--instruct_count={instruct_count}",
-            f"--max_iter={max_iter}"
-        ], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-        
-        stdout, stderr = process.communicate()
-        
-        if process.returncode == 0:
-            subprocess.run(["python", "/app/qwen/collect_data.py"])
-            return "Training completed successfully!"
-        else:
-            return f"Error during training: {stderr.decode()}"
-    except Exception as e:
-        return f"Error: {str(e)}"
+        subprocess.run(["bash", "run.sh"], 
+                      check=True,
+                      cwd="/app/qwen")
+        return "Training completed successfully!"
+    except subprocess.CalledProcessError as e:
+        return f"Error during training: {str(e)}"
 
-# Create Gradio interface
-iface = gr.Interface(
-    fn=run_self_lengthen,
-    inputs=[
-        gr.Textbox(label="Model Path", value="/app/models/base_model"),
-        gr.Number(label="Instruction Count", value=5000, minimum=100),
-        gr.Number(label="Max Iterations", value=3, minimum=1)
-    ],
-    outputs=gr.Textbox(label="Status"),
-    title="Self-Lengthen Training Interface",
-    description="Train language models to generate longer texts using the Self-Lengthen approach."
-)
+# Create the interface
+with gr.Blocks() as iface:
+    gr.Markdown("# Self-Lengthen Training Interface")
+    
+    with gr.Row():
+        with gr.Column():
+            model_path = gr.Textbox(
+                label="Model Path",
+                value="/app/models/base_model",
+                info="Path to the base model"
+            )
+            instruct_count = gr.Number(
+                label="Instruction Count",
+                value=5000,
+                minimum=100,
+                info="Number of instructions to generate"
+            )
+            max_iter = gr.Number(
+                label="Max Iterations",
+                value=3,
+                minimum=1,
+                info="Number of training iterations"
+            )
+            train_btn = gr.Button("Start Training")
+        
+        with gr.Column():
+            status_output = gr.Textbox(
+                label="Status",
+                value="Ready to start training...",
+                interactive=False
+            )
+            refresh_btn = gr.Button("Refresh Status")
+    
+    train_btn.click(
+        fn=start_training,
+        inputs=[model_path, instruct_count, max_iter],
+        outputs=status_output
+    )
+    
+    refresh_btn.click(
+        fn=check_training_status,
+        inputs=None,
+        outputs=status_output
+    )
 
 if __name__ == "__main__":
     iface.launch(server_name="0.0.0.0", server_port=7860)