Update app.py

app.py

@@ -1,123 +1,117 @@
-import os, io, zipfile, shutil, subprocess, json, time, glob, tempfile
-import gradio as gr
+# app.py
+import os, shutil, subprocess, zipfile, time
 from pathlib import Path
-from typing import List, Tuple
+import gradio as gr
 
-WORKDIR      = Path(".")
-DATASET_PATH = WORKDIR / "dataset.jsonl"
-LOG_PATH     = WORKDIR / "train.log"
-MODEL_DIR    = WORKDIR / "trained_model"        # training output folder
-ZIP_PATH     = WORKDIR / "trained_model.zip"    # zipped after train
-MODELS_ROOT  = WORKDIR                          # where we scan for saved AIs
+WORKDIR       = Path(".")
+DATASET_PATH  = WORKDIR / "dataset.jsonl"
+LOG_PATH      = WORKDIR / "train.log"
+MODEL_DIR     = WORKDIR / "trained_model"         # created by train.py
+ZIP_PATH      = WORKDIR / "trained_model.zip"     # created after train
 
 # ---------- helpers ----------
-def _safe_unzip(zip_file: str, out_dir: Path) -> str:
-    out_dir.mkdir(parents=True, exist_ok=True)
-    with zipfile.ZipFile(zip_file, "r") as z:
-        z.extractall(out_dir)
-    # return the inner model folder if zip contained a single directory
-    subdirs = [p for p in out_dir.iterdir() if p.is_dir()]
-    return str(subdirs[0] if len(subdirs) == 1 else out_dir)
-
-def _list_local_models() -> List[str]:
-    """
-    Return model folders found under MODELS_ROOT that look like HF models.
-    We include any folder that has a tokenizer.json or tokenizer_config.json.
-    """
-    candidates = []
-    for p in MODELS_ROOT.iterdir():
-        if not p.is_dir():
-            continue
-        if (p / "tokenizer.json").exists() or (p / "tokenizer_config.json").exists():
-            candidates.append(str(p))
-    return sorted(candidates)
-
-def _start_training_subprocess() -> int:
-    # clear old outputs
+def _list_models():
+    """List model-looking folders in the 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))
+    # also include the default training output if present
+    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/ into trained_model.zip (overwrite if exists)."""
+    if not MODEL_DIR.exists():
+        return False
+    if ZIP_PATH.exists():
+        ZIP_PATH.unlink()
+    shutil.make_archive(ZIP_PATH.with_suffix("").as_posix(), "zip", MODEL_DIR)
+    return ZIP_PATH.exists()
+
+# ---------- TRAIN callbacks ----------
+def upload_dataset(file):
+    if file is None:
+        return "❌ No file selected."
+    shutil.copy(file.name, DATASET_PATH)
+    return f"✅ Uploaded {file.name} → {DATASET_PATH.name}"
+
+def start_training():
+    if not DATASET_PATH.exists():
+        return ("❌ Upload a JSONL first.", "", gr.File.update(visible=False))
+    # clean previous outputs
     if MODEL_DIR.exists():
         shutil.rmtree(MODEL_DIR)
     if ZIP_PATH.exists():
         ZIP_PATH.unlink(missing_ok=True)
+    LOG_PATH.write_text("🔥 Starting training...\n", encoding="utf-8")
 
     cmd = [
         "python", "train.py",
         "--dataset", str(DATASET_PATH),
-        "--output",  str(MODEL_DIR),
-        # sensible defaults for quick, real training; adjust in train.py if needed
+        "--output", str(MODEL_DIR),
         "--model_name", "Salesforce/codegen-350M-multi",
         "--epochs", "1",
         "--batch_size", "2",
         "--block_size", "256",
         "--learning_rate", "5e-5",
+        "--subset", "0"
     ]
-    LOG_PATH.write_text("🔥 Starting training...\n", encoding="utf-8")
     with open(LOG_PATH, "a", encoding="utf-8") as lf:
-        proc = subprocess.Popen(cmd, stdout=lf, stderr=subprocess.STDOUT)
-        return proc.wait()
-
-def _zip_model_folder() -> bool:
-    if not MODEL_DIR.exists():
-        return False
-    if ZIP_PATH.exists():
-        ZIP_PATH.unlink()
-    shutil.make_archive(ZIP_PATH.with_suffix("").as_posix(), "zip", MODEL_DIR)
-    return ZIP_PATH.exists()
-
-# ---------- UI callbacks ----------
-def upload_dataset(file) -> str:
-    if file is None:
-        return "❌ No file selected."
-    shutil.copy(file.name, DATASET_PATH)
-    return f"✅ Uploaded: {file.name} → {DATASET_PATH.name}"
-
-def start_training() -> Tuple[str, str, gr.File]:
-    if not DATASET_PATH.exists():
-        return ("❌ Please upload a JSONL first.", "", gr.File.update(visible=False))
-    exit_code = _start_training_subprocess()
-
-    # after training, try to zip and expose
-    if exit_code == 0 and _zip_model_folder():
-        status = "✅ Training complete."
-        model_info = f"Saved: {MODEL_DIR.name}  |  Zip: {ZIP_PATH.name}"
-        return (status, model_info, gr.File.update(value=str(ZIP_PATH), visible=True))
+        code = subprocess.Popen(cmd, stdout=lf, stderr=subprocess.STDOUT).wait()
+
+    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))
     else:
-        # surface the tail of the log for quick diagnosis
         tail = ""
         if LOG_PATH.exists():
             with open(LOG_PATH, "r", encoding="utf-8") as f:
-                lines = f.readlines()[-30:]
-                tail = "".join(lines)
-        return (f"❌ Training failed (code {exit_code}).", tail, gr.File.update(visible=False))
+                tail = "".join(f.readlines()[-40:])
+        return (f"❌ Training failed (exit {code}). See logs.", tail, gr.File.update(visible=False))
 
-def read_logs() -> str:
+def read_logs():
     if LOG_PATH.exists():
-        return LOG_PATH.read_text(encoding="utf-8")[-20_000:]  # last ~20k chars
-    return "⏳ Waiting for logs..."
+        return LOG_PATH.read_text(encoding="utf-8")[-20000:]
+    return "⏳ Waiting for logs…"
 
-def refresh_model_list() -> List[str]:
-    return _list_local_models()
+def refresh_download():
+    return gr.File.update(value=str(ZIP_PATH), visible=ZIP_PATH.exists())
 
-def upload_model_zip(zip_file) -> Tuple[str, List[str]]:
-    if zip_file is None:
-        return "❌ No zip provided.", refresh_model_list()
-    out = WORKDIR / f"imported_{int(time.time())}"
-    path = _safe_unzip(zip_file.name, out)
-    msg = f"✅ Imported model at: {path}"
-    return msg, refresh_model_list()
+# ---------- TEST callbacks ----------
+def refresh_models():
+    return _list_models()
 
-def generate(model_path: str, prompt: str) -> str:
+def upload_model_zip(zip_file):
+    if zip_file is None:
+        return "❌ No zip selected.", _list_models()
+    # extract to a unique folder
+    dest = WORKDIR / 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()
+
+def generate(model_path, 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)
+        tok = AutoTokenizer.from_pretrained(model_path, use_fast=True)
         if tok.pad_token_id is None and tok.eos_token_id is not None:
             tok.pad_token = tok.eos_token
         model = AutoModelForCausalLM.from_pretrained(model_path)
         gen = pipeline("text-generation", model=model, tokenizer=tok)
 
-        # decoding tuned for code
-        out = gen(
+        text = gen(
             prompt,
             max_new_tokens=220,
             do_sample=True,
@@ -127,45 +121,44 @@ def generate(model_path: str, prompt: str) -> str:
             no_repeat_ngram_size=4,
             eos_token_id=tok.eos_token_id,
             pad_token_id=tok.pad_token_id,
-            truncation=True,
+            truncation=True
         )[0]["generated_text"]
-        return out
+        return text
     except Exception as e:
         return f"❌ Error: {e}"
 
 # ---------- UI ----------
-with gr.Blocks(title="Python AI Trainer") as demo:
-    gr.Markdown("## 🧠 Python AI Trainer\nUpload JSONL, train, then test your model.")
+with gr.Blocks(title="Python AI Trainer") as app:
+    gr.Markdown("## 🧠 Python AI — Train & Test\nUpload JSONL → Train → Download ZIP. Test any stored model separately.")
 
     with gr.Tab("Train"):
-        file_in   = gr.File(label="📥 Upload JSONL Dataset", file_types=[".jsonl", ".jsonl.gz", ".json"])
+        ds = gr.File(label="📥 Upload JSONL dataset", file_types=[".jsonl", ".jsonl.gz", ".json"])
         up_status = gr.Textbox(label="Upload Status", interactive=False)
-        start_btn = gr.Button("🚀 Start Training", variant="primary")
-        logs_box  = gr.Textbox(label="📜 Live Logs (click Refresh)", lines=16)
-        refresh_logs = gr.Button("Refresh Logs")
-        status_box = gr.Textbox(label="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)
-        dl = gr.File(label="📦 Download Trained Model (.zip)", visible=False)
-        refresh_dl = gr.Button("Refresh Download Area")
+        download_file = gr.File(label="📦 Download trained_model.zip", visible=False)
+        refresh_dl_btn = gr.Button("Refresh Download Area")
 
-        file_in.change(fn=upload_dataset, inputs=file_in, outputs=up_status)
-        start_btn.click(fn=start_training, outputs=[status_box, model_info, dl])
-        refresh_logs.click(fn=read_logs, outputs=logs_box)
-        refresh_dl.click(fn=lambda: (gr.File.update(value=str(ZIP_PATH), visible=ZIP_PATH.exists())),
-                         outputs=dl)
+        ds.change(upload_dataset, inputs=ds, outputs=up_status)
+        start.click(start_training, outputs=[status, model_info, download_file])
+        refresh_logs_btn.click(read_logs, outputs=logs)
+        refresh_dl_btn.click(refresh_download, outputs=download_file)
 
     with gr.Tab("Test"):
-        gr.Markdown("### 🔬 Choose a stored AI and prompt it")
-        refresh_models_btn = gr.Button("↻ Refresh AI List")
-        model_list = gr.Dropdown(choices=_list_local_models(), label="Available AIs", interactive=True)
-        up_zip = gr.File(label="Or upload a model .zip to test", file_types=[".zip"])
-        zip_status = gr.Textbox(label="Model Import Status", interactive=False)
-        prompt = gr.Textbox(label="Prompt", lines=6, placeholder="### Instruction:\nPython: write a function ...\n### Response:\n")
-        generate_btn = gr.Button("Generate")
-        output = gr.Textbox(label="AI Response", lines=20)
-
-        refresh_models_btn.click(fn=refresh_model_list, outputs=model_list)
-        up_zip.change(fn=upload_model_zip, inputs=up_zip, outputs=[zip_status, model_list])
-        generate_btn.click(fn=generate, inputs=[model_list, prompt], outputs=output)
-
-demo.launch()
+        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)
+        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])
+        go.click(generate, inputs=[model_list, prompt], outputs=out)
+
+app.launch()