Update app.py

app.py

@@ -3,23 +3,30 @@ import subprocess
 import os
 
 def run_self_lengthen(model_path, instruct_count, max_iter):
-    # Set environment variables
-    os.environ["MODEL_PATH"] = model_path
-    os.environ["INSTRUCT_COUNT"] = str(instruct_count)
-    os.environ["MAX_ITER"] = str(max_iter)
-    
-    # Run the self-lengthen process
     try:
-        subprocess.run(["/app/start.sh"], check=True)
-        return "Training completed successfully! Check the results in the output directory."
-    except subprocess.CalledProcessError as e:
-        return f"Error during training: {str(e)}"
+        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)}"
 
 # Create Gradio interface
 iface = gr.Interface(
     fn=run_self_lengthen,
     inputs=[
-        gr.Textbox(label="Model Path", value="models/base_model"),
+        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)
     ],
@@ -28,4 +35,5 @@ iface = gr.Interface(
     description="Train language models to generate longer texts using the Self-Lengthen approach."
 )
 
-iface.launch(server_name="0.0.0.0", server_port=7860)
+if __name__ == "__main__":
+    iface.launch(server_name="0.0.0.0", server_port=7860)