Update app.py

app.py

@@ -2,40 +2,58 @@ import os, shutil, subprocess, zipfile, time
 from pathlib import Path
 import gradio as gr
 
-ROOT = Path(".").resolve()
-DATASET = ROOT / "dataset.jsonl"
-LOG = ROOT / "train.log"
-OUT_DIR = ROOT / "trained_model"
-ZIP = ROOT / "trained_model.zip"
-PID = ROOT / "TRAIN_PID"
-DONE = ROOT / "TRAIN_DONE"
-ERRF = ROOT / "TRAIN_ERROR"
-
-def ls_workspace():
+ROOT   = Path(_file_).resolve().parent           # /home/user/app
+DATA   = ROOT / "dataset.jsonl"
+LOG    = ROOT / "train.log"
+OUT    = ROOT / "trained_model"
+ZIP    = ROOT / "trained_model.zip"
+DONE   = ROOT / "TRAIN_DONE"                       # <- completion flag
+ERRF   = ROOT / "TRAIN_ERROR"                      # <- error flag
+
+# ---------- helpers ----------
+def ls_workspace() -> str:
     rows = []
     for p in sorted(ROOT.iterdir(), key=lambda x: (x.is_file(), x.name.lower())):
-        sz = p.stat().st_size if p.exists() else 0
-        rows.append(f"{'[DIR]' if p.is_dir() else '     '}\t{sz:>10}\t{p.name}")
+        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 upload_dataset(f):
-    if not f: return "❌ No file.", ls_workspace()
-    shutil.copy(f.name, DATASET)
-    return f"✅ Uploaded → {DATASET.name}", ls_workspace()
+def _reset_artifacts():
+    for path in [OUT, ZIP, DONE, ERRF, LOG]:
+        if path.is_dir():
+            shutil.rmtree(path, ignore_errors=True)
+        else:
+            path.unlink(missing_ok=True)
+
+def _zip_if_ready() -> bool:
+    """Zip OUT → ZIP once DONE exists."""
+    if DONE.exists() and OUT.exists():
+        if ZIP.exists():
+            ZIP.unlink()
+        with zipfile.ZipFile(ZIP, "w", compression=zipfile.ZIP_DEFLATED) as z:
+            for p in OUT.rglob("*"):
+                z.write(p, arcname=p.relative_to(OUT))
+    return ZIP.exists()
+
+# ---------- train tab callbacks ----------
+def upload_dataset(file):
+    if not file:
+        return "❌ No file selected.", ls_workspace()
+    shutil.copy(file.name, DATA)
+    return f"✅ Uploaded → {DATA.name}", ls_workspace()
 
-def start_training():  # non-blocking
-    # clean previous
-    for p in [OUT_DIR, ZIP, DONE, ERRF, PID]:
-        if isinstance(p, Path) and p.is_dir():
-            shutil.rmtree(p, ignore_errors=True)
-        elif isinstance(p, Path) and p.exists():
-            p.unlink(missing_ok=True)
+def start_training():
+    if not DATA.exists():
+        return "❌ Upload a JSONL first.", ls_workspace()
 
+    _reset_artifacts()
     LOG.write_text("🔥 Training started in background…\n", encoding="utf-8")
+
     cmd = [
         "python", "train.py",
-        "--dataset", str(DATASET),
-        "--output",  str(OUT_DIR),
+        "--dataset", str(DATA),
+        "--output",  str(OUT),
         "--model_name", "Salesforce/codegen-350M-multi",
         "--epochs", "1",
         "--batch_size", "2",
@@ -44,32 +62,26 @@ def start_training():  # non-blocking
         "--subset", "0",
     ]
     with open(LOG, "a", encoding="utf-8") as lf:
-        proc = subprocess.Popen(cmd, stdout=lf, stderr=subprocess.STDOUT)
-    PID.write_text(str(proc.pid))
-    return "🚀 Training started. Use “Refresh Logs/Download”.", ls_workspace()
+        subprocess.Popen(cmd, stdout=lf, stderr=subprocess.STDOUT)
+
+    return "🚀 Training launched. Use Refresh buttons.", ls_workspace()
 
 def read_logs():
     return LOG.read_text(encoding="utf-8")[-20000:] if LOG.exists() else "⏳ Waiting…"
 
-def _zip_if_ready():
-    """Zip only when DONE flag exists and zip not created yet."""
-    if DONE.exists() and OUT_DIR.exists() and not ZIP.exists():
-        with zipfile.ZipFile(ZIP, "w", compression=zipfile.ZIP_DEFLATED) as z:
-            for p in OUT_DIR.rglob("*"):
-                z.write(p, arcname=p.relative_to(OUT_DIR))
-    return ZIP.exists()
-
 def refresh_status_and_download():
-    status = "⏳ Training…"
     if ERRF.exists():
-        status = f"❌ Error: {ERRF.read_text(encoding='utf-8')[-500:]}"
+        status = f"❌ Error:\n{ERRF.read_text(encoding='utf-8')[-1200:]}"
     elif DONE.exists():
         status = "✅ Training complete."
+    else:
+        status = "⏳ Training…"
+
     _zip_if_ready()
     files = [str(ZIP)] if ZIP.exists() else []
     return status, gr.Files.update(value=files, visible=bool(files)), ls_workspace()
 
-# ---- Test tab ----
+# ---------- test tab ----------
 def list_models():
     out = []
     for p in ROOT.iterdir():
@@ -77,38 +89,43 @@ def list_models():
             (p / "tokenizer.json").exists() or (p / "tokenizer_config.json").exists()
         ):
             out.append(str(p))
-    if OUT_DIR.exists() and str(OUT_DIR) not in out:
-        out.insert(0, str(OUT_DIR))
+    if OUT.exists() and str(OUT) not in out:
+        out.insert(0, str(OUT))
     return sorted(out)
 
-def import_zip(z):
-    if not z: return "❌ No zip.", list_models()
+def import_zip(zfile):
+    if not zfile:
+        return "❌ No zip selected.", list_models()
     dest = ROOT / f"imported_{int(time.time())}"
     dest.mkdir(parents=True, exist_ok=True)
-    with zipfile.ZipFile(z.name, "r") as zp:
-        zp.extractall(dest)
+    with zipfile.ZipFile(zfile.name, "r") as z:
+        z.extractall(dest)
     return f"✅ Imported to {dest.name}", 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."
+    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)
         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)
-        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"]
+        pipe = pipeline("text-generation", model=model, tokenizer=tok)
+        out = pipe(
+            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 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\nBackground training with reliable zipping.\n")
+    gr.Markdown("## 🧠 Python AI — Train & Test\nBackground training with DONE flag → reliable zip.\n")
 
     with gr.Tab("Train"):
         with gr.Row():