Update app.py

app.py

@@ -1,16 +1,14 @@
-import os, shutil, subprocess, zipfile, time
+# app.py
+import os, shutil, subprocess, zipfile
 from pathlib import Path
 import gradio as gr
 
-ROOT   = Path(__file__).resolve().parent           # /home/user/app
+ROOT   = Path(_file_).resolve().parent           # /home/user/app
 DATA   = ROOT / "dataset.jsonl"
 LOG    = ROOT / "train.log"
 OUT    = ROOT / "trained_model"
 ZIP    = ROOT / "trained_model.zip"
-DONE   = ROOT / "TRAIN_DONE"                       # <- completion flag
-ERRF   = ROOT / "TRAIN_ERROR"                      # <- error flag
 
-# ---------- helpers ----------
 def ls_workspace() -> str:
     rows = []
     for p in sorted(ROOT.iterdir(), key=lambda x: (x.is_file(), x.name.lower())):
@@ -19,24 +17,6 @@ def ls_workspace() -> str:
         rows.append(f"{'[DIR]' if p.is_dir() else '     '}\t{size:>10}\t{p.name}")
     return "\n".join(rows) or "(empty)"
 
-def _reset_artifacts():
-    for path in [OUT, ZIP, DONE, ERRF, LOG]:
-        if path.is_dir():
-            shutil.rmtree(path, ignore_errors=True)
-        else:
-            path.unlink(missing_ok=True)
-
-def _zip_if_ready() -> bool:
-    """Zip OUT → ZIP once DONE exists."""
-    if DONE.exists() and OUT.exists():
-        if ZIP.exists():
-            ZIP.unlink()
-        with zipfile.ZipFile(ZIP, "w", compression=zipfile.ZIP_DEFLATED) as z:
-            for p in OUT.rglob("*"):
-                z.write(p, arcname=p.relative_to(OUT))
-    return ZIP.exists()
-
-# ---------- train tab callbacks ----------
 def upload_dataset(file):
     if not file:
         return "❌ No file selected.", ls_workspace()
@@ -44,44 +24,41 @@ def upload_dataset(file):
     return f"✅ Uploaded → {DATA.name}", ls_workspace()
 
 def start_training():
-    if not DATA.exists():
-        return "❌ Upload a JSONL first.", ls_workspace()
-
-    _reset_artifacts()
-    LOG.write_text("🔥 Training started in background…\n", encoding="utf-8")
+    # clean previous artifacts
+    if OUT.exists(): shutil.rmtree(OUT, ignore_errors=True)
+    if ZIP.exists(): ZIP.unlink(missing_ok=True)
+    LOG.write_text("🔥 Training started…\n", encoding="utf-8")
 
     cmd = [
-        "python", "train.py",
+        "python", str(ROOT / "train.py"),
         "--dataset", str(DATA),
         "--output",  str(OUT),
+        "--zip_path", str(ZIP),
         "--model_name", "Salesforce/codegen-350M-multi",
         "--epochs", "1",
         "--batch_size", "2",
         "--block_size", "256",
         "--learning_rate", "5e-5",
-        "--subset", "0",
     ]
+    # run training (blocking) and capture logs
     with open(LOG, "a", encoding="utf-8") as lf:
-        subprocess.Popen(cmd, stdout=lf, stderr=subprocess.STDOUT)
+        code = subprocess.Popen(cmd, stdout=lf, stderr=subprocess.STDOUT).wait()
 
-    return "🚀 Training launched. Use Refresh buttons.", ls_workspace()
+    # after process exits, show result
+    if code == 0 and ZIP.exists():
+        info = f"✅ Training complete. Saved: {OUT.name} | Zip: {ZIP.name}"
+        return info, gr.File.update(value=str(ZIP), visible=True), ls_workspace(), read_logs()
+    else:
+        info = f"❌ Training failed (exit {code}). See logs."
+        return info, gr.File.update(visible=False), ls_workspace(), read_logs()
 
 def read_logs():
     return LOG.read_text(encoding="utf-8")[-20000:] if LOG.exists() else "⏳ Waiting…"
 
-def refresh_status_and_download():
-    if ERRF.exists():
-        status = f"❌ Error:\n{ERRF.read_text(encoding='utf-8')[-1200:]}"
-    elif DONE.exists():
-        status = "✅ Training complete."
-    else:
-        status = "⏳ Training…"
-
-    _zip_if_ready()
-    files = [str(ZIP)] if ZIP.exists() else []
-    return status, gr.Files.update(value=files, visible=bool(files)), ls_workspace()
+def refresh_download():
+    return gr.File.update(value=str(ZIP), visible=ZIP.exists()), ls_workspace()
 
-# ---------- test tab ----------
+# ---------------- Test tab ----------------
 def list_models():
     out = []
     for p in ROOT.iterdir():
@@ -96,17 +73,16 @@ def list_models():
 def import_zip(zfile):
     if not zfile:
         return "❌ No zip selected.", list_models()
-    dest = ROOT / f"imported_{int(time.time())}"
+    dest = ROOT / f"imported_model"
+    if dest.exists(): shutil.rmtree(dest, ignore_errors=True)
     dest.mkdir(parents=True, exist_ok=True)
     with zipfile.ZipFile(zfile.name, "r") as z:
         z.extractall(dest)
     return f"✅ Imported to {dest.name}", list_models()
 
 def generate(model_path, prompt):
-    if not model_path:
-        return "❌ Select a model."
-    if not prompt or not prompt.strip():
-        return "❌ Enter a prompt."
+    if not model_path: return "❌ Select a model."
+    if not prompt or not prompt.strip(): return "❌ Enter a prompt."
     try:
         from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
         tok = AutoTokenizer.from_pretrained(model_path, use_fast=True)
@@ -123,9 +99,8 @@ def generate(model_path, prompt):
     except Exception as e:
         return f"❌ Error: {e}"
 
-# ---------- UI ----------
 with gr.Blocks(title="Python AI — Train & Test") as app:
-    gr.Markdown("## 🧠 Python AI — Train & Test\nBackground training with DONE flag → reliable zip.\n")
+    gr.Markdown("## 🧠 Python AI — Train & Test (simple + reliable)\nTrainer zips the model itself. UI just shows the zip.\n")
 
     with gr.Tab("Train"):
         with gr.Row():
@@ -133,21 +108,20 @@ with gr.Blocks(title="Python AI — Train & Test") as app:
             ws = gr.Textbox(label="Workspace", lines=16, value=ls_workspace())
         up_status = gr.Textbox(label="Upload Status", interactive=False)
         start = gr.Button("🚀 Start Training", variant="primary")
-        logs = gr.Textbox(label="📜 Logs (click Refresh)", lines=18)
-        refresh_logs_btn = gr.Button("Refresh Logs")
+        logs = gr.Textbox(label="📜 Training Logs", lines=18, value=read_logs())
         status = gr.Textbox(label="Status", interactive=False)
-        downloads = gr.Files(label="📦 Downloads (zips)", value=[], interactive=False)
-        refresh_dl_btn = gr.Button("Refresh Status & Download")
+        download_file = gr.File(label="📦 trained_model.zip", visible=ZIP.exists())
+        refresh_dl_btn = gr.Button("Refresh Download")
 
         ds.change(upload_dataset, inputs=ds, outputs=[up_status, ws])
-        start.click(start_training, outputs=[status, ws])
-        refresh_logs_btn.click(read_logs, outputs=logs)
-        refresh_dl_btn.click(refresh_status_and_download, outputs=[status, downloads, ws])
+        start.click(start_training, outputs=[status, download_file, ws, logs])
+        refresh_dl_btn.click(refresh_download, outputs=[download_file, ws])
 
     with gr.Tab("Test"):
+        gr.Markdown("### Choose a model folder or upload a .zip, then prompt it")
         refresh_btn = gr.Button("↻ Refresh Model List")
         model_list = gr.Dropdown(choices=list_models(), label="Available AIs", interactive=True)
-        zip_in = gr.File(label="Or upload model .zip", file_types=[".zip"])
+        zip_in = gr.File(label="Or upload a model .zip", file_types=[".zip"])
         import_status = gr.Textbox(label="Import Status", interactive=False)
         prompt = gr.Textbox(label="Prompt", lines=8, placeholder="### Instruction:\nPython: write a function ...\n### Response:\n")
         go = gr.Button("Generate")