Update app.py

app.py

@@ -1,25 +1,26 @@
 # app.py
-import os, shutil, subprocess, zipfile, traceback, time, io
+import os, shutil, subprocess, zipfile, traceback, io
 from pathlib import Path
 from datetime import datetime
 import gradio as gr
 
+# ----------------- Paths -----------------
 ROOT = Path(__file__).resolve().parent
 DATA = ROOT / "dataset.jsonl"
 LOG  = ROOT / "train.log"
 RUNS = ROOT / "runs"
 RUNS.mkdir(exist_ok=True)
 
-# -------- logging helpers --------
+# ----------------- Logging -----------------
 def append_log(msg: str):
-    msg = msg.rstrip()
+    msg = (msg or "").rstrip("\n")
     with open(LOG, "a", encoding="utf-8") as lf:
         lf.write(msg + "\n")
 
 def read_logs():
     return LOG.read_text(encoding="utf-8")[-20000:] if LOG.exists() else "⏳ Waiting…"
 
-# -------- workspace + models --------
+# ----------------- Workspace & Models -----------------
 def ls_workspace() -> str:
     rows = []
     for p in sorted(ROOT.iterdir(), key=lambda x: (x.is_file(), x.name.lower())):
@@ -46,7 +47,7 @@ def dropdown_update_safe(models, prefer=None):
     val = prefer if (prefer and prefer in models) else (models[0] if models else None)
     return gr.update(choices=models, value=val)
 
-# -------- dataset upload --------
+# ----------------- Dataset Upload -----------------
 def upload_dataset(file):
     if not file:
         return "❌ No file selected.", ls_workspace()
@@ -55,12 +56,19 @@ def upload_dataset(file):
         return f"✅ Uploaded → {DATA.name}", ls_workspace()
     return "⚠ Unexpected item; please upload a .jsonl file.", ls_workspace()
 
-# -------- training (LIVE LOGS) --------
+# ----------------- Training (Live Logs) -----------------
 def start_training_live(run_name):
     """
     Streams training logs to the UI while the subprocess runs.
     Yields tuples for outputs: [status, download_file, workspace, logs, model_dropdown]
     """
+    # Quick guard: dataset must exist
+    if not DATA.exists():
+        msg = "❌ dataset.jsonl not found. Upload a JSONL dataset first."
+        append_log(msg)
+        yield (msg, gr.update(value=None, visible=False), ls_workspace(), read_logs(), dropdown_update_safe(list_models()))
+        return
+
     run_id = (run_name or "").strip() or datetime.now().strftime("run_%Y%m%d_%H%M%S")
     out_dir  = RUNS / run_id
     zip_path = RUNS / f"{run_id}.zip"
@@ -89,18 +97,25 @@ def start_training_live(run_name):
     append_log("▶ " + " ".join(cmd))
 
     # start subprocess with live stdout
-    proc = subprocess.Popen(
-        cmd,
-        stdout=subprocess.PIPE,
-        stderr=subprocess.STDOUT,
-        bufsize=1,
-        universal_newlines=True,
-        encoding="utf-8",
-        errors="replace",
-    )
+    try:
+        proc = subprocess.Popen(
+            cmd,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            bufsize=1,
+            universal_newlines=True,
+            encoding="utf-8",
+            errors="replace",
+        )
+    except Exception as e:
+        err = "❌ Failed to start train.py: " + "".join(traceback.format_exception_only(type(e), e))
+        append_log(err)
+        yield (err, gr.update(value=None, visible=False), ls_workspace(), read_logs(), dropdown_update_safe(list_models()))
+        return
 
     live_log = io.StringIO()
     status_msg = f"🚀 Training run '{run_id}' in progress…"
+
     # stream loop
     while True:
         line = proc.stdout.readline()
@@ -154,7 +169,7 @@ def refresh_download():
         dropdown_update_safe(models)
     )
 
-# -------- import a zip as a model folder --------
+# ----------------- Import a Zip as Model Folder -----------------
 def import_zip(zfile):
     if not zfile:
         return "❌ No zip selected.", list_models()
@@ -166,7 +181,7 @@ def import_zip(zfile):
         z.extractall(dest)
     return f"✅ Imported to {dest.name}", list_models()
 
-# -------- generation: cached pipeline --------
+# ----------------- Generation (cached pipeline) -----------------
 _GEN_CACHE = {"path": None, "pipe": None}
 
 def get_generation_pipeline(model_path: str):
@@ -201,52 +216,93 @@ def get_generation_pipeline(model_path: str):
     append_log("✅ Pipeline loaded.")
     return pipe
 
-def generate(model_path, prompt):
-    from pathlib import Path
+# ----------------- Test Tab Helpers -----------------
+def ping():
+    append_log("🔔 Ping pressed (UI wiring OK)")
+    return "✅ UI is connected and responding."
+
+def load_selected_model(model_path):
+    # Dropdown may pass a list; coerce to string
+    if isinstance(model_path, list):
+        model_path = model_path[0] if model_path else None
+    if not model_path:
+        return "❌ Select a model first."
+    if not isinstance(model_path, str):
+        return f"❌ Invalid model path type: {type(model_path)._name_}"
+    p = Path(model_path)
+    if not p.exists() or not p.is_dir():
+        return f"❌ Model folder not found: {model_path}"
+    try:
+        append_log(f"📦 Load request → {model_path}")
+        _ = get_generation_pipeline(model_path)
+        append_log(f"✅ Loaded pipeline: {model_path}")
+        return f"✅ Loaded: {model_path}"
+    except Exception as e:
+        tb = traceback.format_exc()
+        append_log("❌ Load error:\n" + tb)
+        return "❌ Error while loading model:\n" + "".join(traceback.format_exception_only(type(e), e))
+
+def generate_stream(model_path, prompt):
+    """Stream intermediate status to prove the button is working, then final text."""
+    # immediate feedback
+    yield "⏳ Loading model…"
+    append_log("▶ Generate clicked")
 
-    # Coerce Dropdown value (can be list)
+    # Coerce
     if isinstance(model_path, list):
         model_path = model_path[0] if model_path else None
 
     # validate
     if not model_path:
-        return "❌ Select a model from the dropdown first."
+        msg = "❌ Select a model from the dropdown first."
+        append_log(msg); yield msg; return
     if not isinstance(model_path, str):
-        # fix minor bug: _name_ not name
-        return f"❌ Invalid model path type: {type(model_path)._name_}"
+        msg = f"❌ Invalid model path type: {type(model_path)._name_}"
+        append_log(msg); yield msg; return
     if not Path(model_path).exists():
-        return f"❌ Model folder not found: {model_path}"
+        msg = f"❌ Model folder not found: {model_path}"
+        append_log(msg); yield msg; return
     if not prompt or not prompt.strip():
-        return "❌ Enter a prompt."
+        msg = "❌ Enter a prompt."
+        append_log(msg); yield msg; return
 
     try:
         pipe = get_generation_pipeline(model_path)
-        append_log(f"📝 Generating for prompt ({len(prompt)} chars)…")
-        out = pipe(
+        yield "⚙ Generating… (this may take a bit on CPU)"
+        append_log(f"📝 Generating… prompt_len={len(prompt)}")
+        result = pipe(
             prompt.strip(),
-            max_new_tokens=120,
+            max_new_tokens=80,      # quicker to show something
             do_sample=True,
-            temperature=0.4,
+            temperature=0.3,
             top_p=0.9,
             repetition_penalty=1.15,
             no_repeat_ngram_size=4,
-            truncation=True
-        )[0]["generated_text"]
+            truncation=True,
+            return_full_text=True,
+        )
+        text = result[0].get("generated_text", "")
+        if not text:
+            append_log("⚠ Empty generated_text")
+            yield "⚠ Model returned empty text. Try lowering temperature or adding more context."
+            return
         append_log("✅ Generation OK.")
-        return out
+        yield text
     except Exception as e:
         tb = traceback.format_exc()
         append_log("❌ Generation error:\n" + tb)
-        return "❌ Error during generation:\n" + "".join(traceback.format_exception_only(type(e), e))
+        yield "❌ Error during generation:\n" + "".join(traceback.format_exception_only(type(e), e))
 
-# -------- UI --------
+# ----------------- UI -----------------
 with gr.Blocks(title="Python AI — Train & Test") as app:
     gr.Markdown("## 🧠 Python AI — Train & Test\n• Unique runs • Safe download • Cached generation • Live logs\n")
 
-    # Test first (so Train can update its dropdown)
+    # ---------- Test Tab ----------
     with gr.Tab("Test"):
         gr.Markdown("### Choose a model folder or upload a .zip, then prompt it")
-        refresh_btn = gr.Button("↻ Refresh Model List")
+        with gr.Row():
+            refresh_btn = gr.Button("↻ Refresh Model List")
+            ping_btn = gr.Button("🔔 Ping UI")  # quick sanity check
         model_list = gr.Dropdown(
             choices=list_models(),
             label="Available AIs",
@@ -254,12 +310,21 @@ with gr.Blocks(title="Python AI — Train & Test") as app:
             allow_custom_value=True,  # keeps UI quiet when empty
             multiselect=False         # force single selection
         )
+        load_btn = gr.Button("📦 Load Model")
+        load_status = gr.Textbox(label="Model Status", interactive=False)
+
         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")
+
+        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 ----------
     with gr.Tab("Train"):
         with gr.Row():
             ds = gr.File(label="📥 Upload JSONL", file_types=[".jsonl"])
@@ -272,22 +337,30 @@ with gr.Blocks(title="Python AI — Train & Test") as app:
         download_file = gr.File(label="📦 Latest trained zip", visible=False)
         refresh_dl_btn = gr.Button("Refresh Download")
 
-    # wiring
+    # ---------- Wiring ----------
+    # Upload + workspace
     ds.change(upload_dataset, inputs=ds, outputs=[up_status, ws])
 
-    # STREAMED training: function yields updates
+    # Train (live streaming)
     start.click(
         start_training_live,
         inputs=[run_name],
         outputs=[status, download_file, ws, logs, model_list]
     )
 
+    # Download refresh
     refresh_dl_btn.click(
         refresh_download,
         outputs=[download_file, ws, model_list]
     )
+
+    # Test tab helpers
     refresh_btn.click(lambda: dropdown_update_safe(list_models()), outputs=model_list)
+    ping_btn.click(ping, outputs=out)
+    load_btn.click(load_selected_model, inputs=[model_list], outputs=[load_status])
     zip_in.change(import_zip, inputs=zip_in, outputs=[import_status, model_list])
-    go.click(generate, inputs=[model_list, prompt], outputs=out)
+
+    # Streamed generation output
+    go.click(generate_stream, inputs=[model_list, prompt], outputs=out)
 
 app.queue(default_concurrency_limit=1).launch()