Update app.py

app.py

@@ -3,7 +3,7 @@ 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
 DATA   = ROOT / "dataset.jsonl"
 LOG    = ROOT / "train.log"
 OUT    = ROOT / "trained_model"
@@ -17,6 +17,17 @@ 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 list_models():
+    out = []
+    for p in ROOT.iterdir():
+        if p.is_dir() and (p / "config.json").exists() and (
+            (p / "tokenizer.json").exists() or (p / "tokenizer_config.json").exists()
+        ):
+            out.append(str(p))
+    if OUT.exists() and str(OUT) not in out:
+        out.insert(0, str(OUT))
+    return sorted(out)
+
 def upload_dataset(file):
     if not file:
         return "❌ No file selected.", ls_workspace()
@@ -43,35 +54,29 @@ def start_training():
     with open(LOG, "a", encoding="utf-8") as lf:
         code = subprocess.Popen(cmd, stdout=lf, stderr=subprocess.STDOUT).wait()
 
+    # build model-list update payload
+    models = list_models()
+    model_update = gr.update(choices=models, value=str(OUT) if OUT.exists() else None)
+
     if code == 0 and ZIP.exists():
         info = f"✅ Training complete. Saved: {OUT.name} | Zip: {ZIP.name}"
-        return info, gr.update(value=str(ZIP), visible=True), ls_workspace(), read_logs()
+        return info, gr.update(value=str(ZIP), visible=True), ls_workspace(), read_logs(), model_update
     else:
         info = f"❌ Training failed (exit {code}). See logs."
-        return info, gr.update(visible=False), ls_workspace(), read_logs()
+        return info, gr.update(visible=False), ls_workspace(), read_logs(), model_update
 
 def read_logs():
     return LOG.read_text(encoding="utf-8")[-20000:] if LOG.exists() else "⏳ Waiting…"
 
 def refresh_download():
-    return gr.update(value=str(ZIP), visible=ZIP.exists()), ls_workspace()
-
-# ---------------- Test tab ----------------
-def list_models():
-    out = []
-    for p in ROOT.iterdir():
-        if p.is_dir() and (p / "config.json").exists() and (
-            (p / "tokenizer.json").exists() or (p / "tokenizer_config.json").exists()
-        ):
-            out.append(str(p))
-    if OUT.exists() and str(OUT) not in out:
-        out.insert(0, str(OUT))
-    return sorted(out)
+    # also refresh model dropdown
+    models = list_models()
+    return gr.update(value=str(ZIP), visible=ZIP.exists()), ls_workspace(), gr.update(choices=models)
 
 def import_zip(zfile):
     if not zfile:
         return "❌ No zip selected.", list_models()
-    dest = ROOT / f"imported_model"
+    dest = ROOT / "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:
@@ -98,8 +103,20 @@ def generate(model_path, prompt):
         return f"❌ Error: {e}"
 
 with gr.Blocks(title="Python AI — Train & Test") as app:
-    gr.Markdown("## 🧠 Python AI — Train & Test (simple + reliable)\nTrainer zips the model itself. UI just shows the zip.\n")
+    gr.Markdown("## 🧠 Python AI — Train & Test (auto-add to Test tab)")
+
+    # --- Test tab UI FIRST so we can reference components ---
+    with gr.Tab("Test"):
+        gr.Markdown("### Pick a model or upload a .zip")
+        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 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")
+        out = gr.Textbox(label="AI Response", lines=20)
 
+    # --- Train tab UI ---
     with gr.Tab("Train"):
         with gr.Row():
             ds = gr.File(label="📥 Upload JSONL", file_types=[".jsonl", ".jsonl.gz", ".json"])
@@ -111,22 +128,18 @@ with gr.Blocks(title="Python AI — Train & Test") as app:
         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, 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 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")
-        out = gr.Textbox(label="AI Response", lines=20)
-
-        refresh_btn.click(list_models, outputs=model_list)
-        zip_in.change(import_zip, inputs=zip_in, outputs=[import_status, model_list])
-        go.click(generate, inputs=[model_list, prompt], outputs=out)
+    # --- wiring ---
+    ds.change(upload_dataset, inputs=ds, outputs=[up_status, ws])
+    start.click(
+        start_training,
+        outputs=[status, download_file, ws, logs, model_list]  # <-- update Test dropdown automatically
+    )
+    refresh_dl_btn.click(
+        refresh_download,
+        outputs=[download_file, ws, model_list]               # <-- also updates Test dropdown
+    )
+    refresh_btn.click(lambda: gr.update(choices=list_models()), outputs=model_list)
+    zip_in.change(import_zip, inputs=zip_in, outputs=[import_status, model_list])
+    go.click(generate, inputs=[model_list, prompt], outputs=out)
 
 app.launch()