Update app.py

app.py

@@ -3,28 +3,28 @@ import os, shutil, subprocess, zipfile, time
 from pathlib import Path
 import gradio as gr
 
-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
+# --- paths ---
+WORKDIR      = Path(".")
+DATASET_PATH = WORKDIR / "dataset.jsonl"
+LOG_PATH     = WORKDIR / "train.log"
+MODEL_DIR    = WORKDIR / "trained_model"
+ZIP_PATH     = WORKDIR / "trained_model.zip"
 
 # ---------- helpers ----------
 def _list_models():
-    """List model-looking folders in the workspace."""
+    """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))
-    # 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)."""
+    """Zip trained_model/ -> trained_model.zip"""
     if not MODEL_DIR.exists():
         return False
     if ZIP_PATH.exists():
@@ -42,7 +42,8 @@ def upload_dataset(file):
 def start_training():
     if not DATASET_PATH.exists():
         return ("❌ Upload a JSONL first.", "", gr.File.update(visible=False))
-    # clean previous outputs
+
+    # clean previous artifacts
     if MODEL_DIR.exists():
         shutil.rmtree(MODEL_DIR)
     if ZIP_PATH.exists():
@@ -52,13 +53,13 @@ def start_training():
     cmd = [
         "python", "train.py",
         "--dataset", str(DATASET_PATH),
-        "--output", str(MODEL_DIR),
+        "--output",  str(MODEL_DIR),
         "--model_name", "Salesforce/codegen-350M-multi",
         "--epochs", "1",
         "--batch_size", "2",
         "--block_size", "256",
         "--learning_rate", "5e-5",
-        "--subset", "0"
+        "--subset", "0",
     ]
     with open(LOG_PATH, "a", encoding="utf-8") as lf:
         code = subprocess.Popen(cmd, stdout=lf, stderr=subprocess.STDOUT).wait()
@@ -91,7 +92,6 @@ def refresh_models():
 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:
@@ -128,7 +128,7 @@ def generate(model_path, prompt):
         return f"❌ Error: {e}"
 
 # ---------- UI ----------
-with gr.Blocks(title="Python AI Trainer") as app:
+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.")
 
     with gr.Tab("Train"):