Update app.py

app.py

@@ -1,9 +1,10 @@
+# app.py
 import os, shutil, subprocess, zipfile
 from pathlib import Path
 import gradio as gr
 
 ROOT   = Path(__file__).resolve().parent
-DATA   = ROOT / "dataset.jsonl"             # single-file mode target
+DATA   = ROOT / "dataset.jsonl"
 LOG    = ROOT / "train.log"
 OUT    = ROOT / "trained_model"
 ZIP    = ROOT / "trained_model.zip"
@@ -30,15 +31,15 @@ def list_models():
         out.insert(0, str(OUT))
     return sorted(out)
 
+def dropdown_update_safe(models, prefer=None):
+    """Return a gr.update that always uses a value present in choices (or None)."""
+    val = prefer if (prefer and prefer in models) else (models[0] if models else None)
+    return gr.update(choices=models, value=val)
+
 # ---------- train tab ----------
 def upload_dataset(file):
-    """
-    If user uploads a file -> copy to dataset.jsonl
-    If user uploads a folder -> we DO NOT move it, they’ll pass folder path via a textbox if needed.
-    """
     if not file:
         return "❌ No file selected.", ls_workspace()
-    # If it's a file object, copy to DATA
     if hasattr(file, "name") and os.path.isfile(file.name):
         shutil.copy(file.name, DATA)
         return f"✅ Uploaded → {DATA.name}", ls_workspace()
@@ -52,10 +53,9 @@ def start_training():
         ZIP.unlink(missing_ok=True)
     LOG.write_text("🔥 Training started…\n", encoding="utf-8")
 
-    # Run trainer (blocking) and capture output in train.log
     cmd = [
         "python", str(ROOT / "train.py"),
-        "--dataset", str(DATA),                  # For folder-mode, replace DATA with folder path in train.py if you extend UI
+        "--dataset", str(DATA),
         "--output",  str(OUT),
         "--zip_path", str(ZIP),
         "--model_name", "Salesforce/codegen-350M-multi",
@@ -67,10 +67,9 @@ def start_training():
     with open(LOG, "a", encoding="utf-8") as lf:
         code = subprocess.Popen(cmd, stdout=lf, stderr=subprocess.STDOUT).wait()
 
-    # Refresh model list & set selection only if it’s present
     models = list_models()
-    selected = str(OUT) if OUT.exists() and str(OUT) in models else None
-    model_update = gr.update(choices=models, value=selected)
+    prefer = str(OUT) if OUT.exists() else None
+    model_update = dropdown_update_safe(models, prefer=prefer)
 
     if code == 0 and ZIP.exists():
         info = f"✅ Training complete. Saved: {OUT.name} | Zip: {ZIP.name}"
@@ -84,7 +83,11 @@ def read_logs():
 
 def refresh_download():
     models = list_models()
-    return gr.update(value=(str(ZIP) if ZIP.exists() else None), visible=ZIP.exists()), ls_workspace(), gr.update(choices=models)
+    return (
+        gr.update(value=(str(ZIP) if ZIP.exists() else None), visible=ZIP.exists()),
+        ls_workspace(),
+        dropdown_update_safe(models)
+    )
 
 # ---------- test tab ----------
 def import_zip(zfile):
@@ -125,11 +128,16 @@ def generate(model_path, prompt):
 with gr.Blocks(title="Python AI — Train & Test") as app:
     gr.Markdown("## 🧠 Python AI — Train & Test\nTrainer saves & zips. UI only shows existing artifacts.\n")
 
-    # Test tab (declared first so we can update its dropdown from Train tab)
+    # Test first (so Train can update it)
     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)
+        model_list = gr.Dropdown(
+            choices=list_models(),
+            label="Available AIs",
+            interactive=True,
+            allow_custom_value=True  # <-- stops warnings when choices are empty
+        )
         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")
@@ -159,7 +167,7 @@ with gr.Blocks(title="Python AI — Train & Test") as app:
         outputs=[download_file, ws, model_list]
     )
 
-    refresh_btn.click(lambda: gr.update(choices=list_models()), outputs=model_list)
+    refresh_btn.click(lambda: dropdown_update_safe(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)