Update app.py

app.py

@@ -3,47 +3,53 @@ import os, shutil, subprocess, zipfile, time
 from pathlib import Path
 import gradio as gr
 
-# --- paths ---
-WORKDIR      = Path(".")
-DATASET_PATH = WORKDIR / "dataset.jsonl"
-LOG_PATH     = WORKDIR / "train.log"
-MODEL_DIR    = WORKDIR / "trained_model"
-ZIP_PATH     = WORKDIR / "trained_model.zip"
+ROOT         = Path(".").resolve()
+DATASET_PATH = ROOT / "dataset.jsonl"
+LOG_PATH     = ROOT / "train.log"
+MODEL_DIR    = ROOT / "trained_model"
+ZIP_PATH     = ROOT / "trained_model.zip"
 
 # ---------- helpers ----------
-def _list_models():
-    """List model-like folders in workspace."""
-    out = []
-    for p in WORKDIR.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 MODEL_DIR.exists() and str(MODEL_DIR) not in out:
-        out.insert(0, str(MODEL_DIR))
-    return sorted(out)
-
-def _zip_model_folder():
-    """Zip trained_model/ -> trained_model.zip"""
+def list_workspace():
+    rows = []
+    for p in sorted(ROOT.iterdir(), key=lambda x: (x.is_file(), x.name.lower())):
+        try:
+            size = p.stat().st_size
+        except Exception:
+            size = 0
+        rows.append(f"{'[DIR]' if p.is_dir() else '     '}\t{size:>10}\t{p.name}")
+    return "\n".join(rows) or "(empty)"
+
+def list_zips():
+    return [str(p) for p in ROOT.glob("*.zip")]
+
+def zip_trained_model():
     if not MODEL_DIR.exists():
-        return False
+        return False, "trained_model/ not found"
+    # remove old zip
     if ZIP_PATH.exists():
-        ZIP_PATH.unlink()
-    shutil.make_archive(ZIP_PATH.with_suffix("").as_posix(), "zip", MODEL_DIR)
-    return ZIP_PATH.exists()
+        try:
+            ZIP_PATH.unlink()
+        except Exception as e:
+            return False, f"could not remove old zip: {e}"
+    # create zip
+    try:
+        with zipfile.ZipFile(ZIP_PATH, "w", compression=zipfile.ZIP_DEFLATED) as z:
+            for path in MODEL_DIR.rglob("*"):
+                z.write(path, arcname=path.relative_to(MODEL_DIR))
+    except Exception as e:
+        return False, f"zip error: {e}"
+    return ZIP_PATH.exists(), f"created {ZIP_PATH.name}"
 
-# ---------- TRAIN callbacks ----------
+# ---------- train ----------
 def upload_dataset(file):
-    if file is None:
-        return "❌ No file selected."
+    if not file:
+        return "❌ No file selected.", list_workspace()
     shutil.copy(file.name, DATASET_PATH)
-    return f"✅ Uploaded {file.name} → {DATASET_PATH.name}"
+    return f"✅ Uploaded → {DATASET_PATH.name}", list_workspace()
 
 def start_training():
-    if not DATASET_PATH.exists():
-        return ("❌ Upload a JSONL first.", "", gr.File.update(visible=False))
-
-    # clean previous artifacts
+    # clean
     if MODEL_DIR.exists():
         shutil.rmtree(MODEL_DIR)
     if ZIP_PATH.exists():
@@ -64,18 +70,18 @@ def start_training():
     with open(LOG_PATH, "a", encoding="utf-8") as lf:
         code = subprocess.Popen(cmd, stdout=lf, stderr=subprocess.STDOUT).wait()
 
+    # zip if success
     if code == 0:
-        ok = _zip_model_folder()
-        info = f"Saved to: {MODEL_DIR.name}"
-        if ok:
-            info += f" | Zip: {ZIP_PATH.name}"
-        return ("✅ Training complete.", info, gr.File.update(value=str(ZIP_PATH), visible=ok))
+        ok, msg = zip_trained_model()
+        info = f"Saved to: {MODEL_DIR.name} | {msg}"
+        files = list_zips() if ok else []
+        return ("✅ Training complete.", info, gr.Files.update(value=files, visible=ok), list_workspace())
     else:
         tail = ""
         if LOG_PATH.exists():
             with open(LOG_PATH, "r", encoding="utf-8") as f:
-                tail = "".join(f.readlines()[-40:])
-        return (f"❌ Training failed (exit {code}). See logs.", tail, gr.File.update(visible=False))
+                tail = "".join(f.readlines()[-60:])
+        return (f"❌ Training failed (exit {code}). See logs below.", tail, gr.Files.update(visible=False), list_workspace())
 
 def read_logs():
     if LOG_PATH.exists():
@@ -83,20 +89,30 @@ def read_logs():
     return "⏳ Waiting for logs…"
 
 def refresh_download():
-    return gr.File.update(value=str(ZIP_PATH), visible=ZIP_PATH.exists())
+    files = list_zips()
+    return gr.Files.update(value=files, visible=bool(files)), list_workspace()
 
-# ---------- TEST callbacks ----------
-def refresh_models():
-    return _list_models()
+# ---------- test ----------
+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))
+    # ensure trained_model on top if present
+    if MODEL_DIR.exists() and str(MODEL_DIR) not in out:
+        out.insert(0, str(MODEL_DIR))
+    return sorted(out)
 
-def upload_model_zip(zip_file):
-    if zip_file is None:
-        return "❌ No zip selected.", _list_models()
-    dest = WORKDIR / f"imported_{int(time.time())}"
+def import_zip(zip_file):
+    if not zip_file:
+        return "❌ No zip selected.", list_models()
+    dest = ROOT / f"imported_{int(time.time())}"
     dest.mkdir(parents=True, exist_ok=True)
     with zipfile.ZipFile(zip_file.name, "r") as z:
         z.extractall(dest)
-    return f"✅ Imported to {dest}", _list_models()
+    return f"✅ Imported to {dest.name}", list_models()
 
 def generate(model_path, prompt):
     if not model_path:
@@ -110,55 +126,48 @@ def generate(model_path, prompt):
             tok.pad_token = tok.eos_token
         model = AutoModelForCausalLM.from_pretrained(model_path)
         gen = pipeline("text-generation", model=model, tokenizer=tok)
-
-        text = gen(
-            prompt,
-            max_new_tokens=220,
-            do_sample=True,
-            temperature=0.2,
-            top_p=0.9,
-            repetition_penalty=1.2,
-            no_repeat_ngram_size=4,
-            eos_token_id=tok.eos_token_id,
-            pad_token_id=tok.pad_token_id,
-            truncation=True
+        out = gen(
+            prompt, max_new_tokens=220, do_sample=True, temperature=0.2, top_p=0.9,
+            repetition_penalty=1.2, no_repeat_ngram_size=4,
+            eos_token_id=tok.eos_token_id, pad_token_id=tok.pad_token_id, truncation=True
         )[0]["generated_text"]
-        return text
+        return out
     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\nUpload JSONL → Train → Download ZIP. Test any stored model separately.")
+    gr.Markdown("## 🧠 Python AI — Train & Test\nTrain → Zip → Download. Test models separately.\n")
 
     with gr.Tab("Train"):
-        ds = gr.File(label="📥 Upload JSONL dataset", file_types=[".jsonl", ".jsonl.gz", ".json"])
+        with gr.Row():
+            ds = gr.File(label="📥 Upload JSONL dataset", file_types=[".jsonl", ".jsonl.gz", ".json"])
+            ws = gr.Textbox(label="Workspace Explorer", lines=16, value=list_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")
         status = gr.Textbox(label="Status", interactive=False)
         model_info = gr.Textbox(label="Model Output", interactive=False)
-        download_file = gr.File(label="📦 Download trained_model.zip", visible=False)
+        downloads = gr.Files(label="📦 Downloads (zips)", value=list_zips(), interactive=False)
         refresh_dl_btn = gr.Button("Refresh Download Area")
 
-        ds.change(upload_dataset, inputs=ds, outputs=up_status)
-        start.click(start_training, outputs=[status, model_info, download_file])
+        ds.change(upload_dataset, inputs=ds, outputs=[up_status, ws])
+        start.click(start_training, outputs=[status, model_info, downloads, ws])
         refresh_logs_btn.click(read_logs, outputs=logs)
-        refresh_dl_btn.click(refresh_download, outputs=download_file)
+        refresh_dl_btn.click(refresh_download, outputs=[downloads, ws])
 
     with gr.Tab("Test"):
-        gr.Markdown("### 🔬 Pick a stored AI (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)
         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(refresh_models, outputs=model_list)
-        zip_in.change(upload_model_zip, inputs=zip_in, outputs=[import_status, model_list])
+        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)
 
 app.launch()