Update app.py

app.py

@@ -10,128 +10,148 @@ import glob
 import gradio as gr
 from transformers import pipeline
 
-# ---- constants / paths ----
 LOG_FILE = "train.log"
+GEN_LOG_FILE = "dataset_gen.log"
 MODEL_DIR = "trained_model"
 ZIP_FILE = "trained_model.zip"
-ZIP_TEMP = ZIP_FILE + ".part"  # atomic write to avoid half-written zips
+ZIP_TEMP = ZIP_FILE + ".part"
 
-# ---- utils ----
 def _human_size(nbytes: int) -> str:
-    units = ["B", "KB", "MB", "GB", "TB"]
-    i, x = 0, float(nbytes)
-    while x >= 1024 and i < len(units) - 1:
-        x /= 1024.0
-        i += 1
+    units = ["B","KB","MB","GB","TB"]; i=0; x=float(nbytes)
+    while x>=1024 and i<len(units)-1: x/=1024.0; i+=1
     return f"{x:.1f} {units[i]}"
 
 def _read_file_safely(path: str, fallback: str):
     if os.path.exists(path):
         try:
-            with open(path, "r", encoding="utf-8", errors="ignore") as f:
-                return f.read()
-        except:
-            return fallback
+            with open(path,"r",encoding="utf-8",errors="ignore") as f: return f.read()
+        except: return fallback
     return fallback
 
 def _zip_folder_atomic(src_dir: str, zip_path: str, tmp_path: str):
-    if os.path.exists(tmp_path):
-        os.remove(tmp_path)
-    with zipfile.ZipFile(tmp_path, "w", compression=zipfile.ZIP_DEFLATED) as zf:
-        for root, _, files in os.walk(src_dir):
+    if os.path.exists(tmp_path): os.remove(tmp_path)
+    with zipfile.ZipFile(tmp_path,"w",compression=zipfile.ZIP_DEFLATED) as zf:
+        for root,_,files in os.walk(src_dir):
             for fn in files:
-                full = os.path.join(root, fn)
-                arc = os.path.relpath(full, src_dir)
-                zf.write(full, arcname=arc)
-    if os.path.exists(zip_path):
-        os.remove(zip_path)
-    os.replace(tmp_path, zip_path)
+                full=os.path.join(root,fn); arc=os.path.relpath(full,src_dir)
+                zf.write(full,arcname=arc)
+    if os.path.exists(zip_path): os.remove(zip_path)
+    os.replace(tmp_path,zip_path)
 
 def _download_info_text() -> str:
-    if not os.path.exists(ZIP_FILE):
-        return "No trained model yet."
-    size = _human_size(os.path.getsize(ZIP_FILE))
-    mtime = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(os.path.getmtime(ZIP_FILE)))
+    if not os.path.exists(ZIP_FILE): return "No trained model yet."
+    size=_human_size(os.path.getsize(ZIP_FILE))
+    mtime=time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(os.path.getmtime(ZIP_FILE)))
     return f"*Model ready:* {ZIP_FILE}  \n*Size:* {size}  \n*Last modified:* {mtime}"
 
 def ensure_clean_zip():
     for p in (ZIP_FILE, ZIP_TEMP):
         if os.path.exists(p):
+            try: os.remove(p)
+            except: pass
+
+# --------- Dataset Generator ----------
+def start_generation(total, shard_size, out_dir, prefix):
+    total=int(total or 1_000_000)
+    shard_size=int(shard_size or 10_000)
+    out_dir=(out_dir or "json_dataset_v1").strip()
+    prefix=(prefix or "json").strip()
+    with open(GEN_LOG_FILE,"w") as log:
+        log.write(f"🚧 Generating dataset: total={total}, shard_size={shard_size}, out_dir={out_dir}, prefix={prefix}\n")
+    def _worker():
+        with open(GEN_LOG_FILE,"a") as log:
+            if not os.path.exists("make_json_dataset.py"):
+                log.write("❌ make_json_dataset.py not found.\n"); return
             try:
-                os.remove(p)
-            except:
-                pass
+                p = subprocess.Popen(
+                    ["python","make_json_dataset.py",
+                     "--total",str(total),
+                     "--shard_size",str(shard_size),
+                     "--out_dir",out_dir,
+                     "--prefix",prefix],
+                    stdout=log, stderr=subprocess.STDOUT
+                )
+                p.wait()
+                log.write(f"\n🔚 Generator exited with code {p.returncode}\n")
+                if p.returncode==0:
+                    files = sorted(glob.glob(os.path.join(out_dir,"*.jsonl.gz")))
+                    log.write(f"✅ Done. Shards: {len(files)} in {out_dir}\n")
+                else:
+                    log.write("❌ Generation failed.\n")
+            except Exception as e:
+                log.write(f"\n❌ Exception: {e}\n")
+    threading.Thread(target=_worker, daemon=True).start()
+    return f"🚀 Dataset generation started. Output folder: {out_dir}"
+
+def read_gen_logs():
+    return _read_file_safely(GEN_LOG_FILE,"Waiting for generator logs...")
 
-# ---- training helpers ----
+def list_shards(folder):
+    if not folder or not os.path.isdir(folder): return "❌ Provide a valid folder path."
+    jsonl = sorted(glob.glob(os.path.join(folder,"*.jsonl")))
+    gz    = sorted(glob.glob(os.path.join(folder,"*.jsonl.gz")))
+    total = len(jsonl)+len(gz)
+    if total==0: return "No shards found."
+    preview=(jsonl+gz)[:10]
+    lines=[f"Found {total} shard(s). Showing first {len(preview)}:"]+[f"- {os.path.basename(p)}" for p in preview]
+    return "\n".join(lines)
+
+# --------- Training ----------
 def upload_file(file):
-    """Save uploaded dataset to uploads/ and return stable path."""
-    if file is None:
-        return "❌ No file uploaded.", ""
+    if file is None: return "❌ No file uploaded.", ""
     os.makedirs("uploads", exist_ok=True)
     dst = os.path.join("uploads", f"dataset_{uuid.uuid4().hex}.jsonl")
     shutil.copy(file.name, dst)
     return f"✅ Uploaded: {os.path.basename(file.name)} → {dst}", dst
 
 def _train_single_file(dataset_path: str, log):
-    """Call train.py once for a single json/jsonl (or .gz expanded) file."""
-    proc = subprocess.Popen(
-        ["python", "train.py", "--dataset", dataset_path, "--output", MODEL_DIR],
-        stdout=log,
-        stderr=subprocess.STDOUT,
-    )
-    proc.wait()
-    log.write(f"\n    ↳ train.py exited {proc.returncode} for {os.path.basename(dataset_path)}\n")
-    return proc.returncode == 0
+    p = subprocess.Popen(["python","train.py","--dataset",dataset_path,"--output",MODEL_DIR],
+                         stdout=log, stderr=subprocess.STDOUT)
+    p.wait()
+    log.write(f"\n    ↳ train.py exited {p.returncode} for {os.path.basename(dataset_path)}\n")
+    return p.returncode==0
 
 def _train_worker(dataset_path: str, shards_folder: str):
-    with open(LOG_FILE, "w") as log:
-        log.write("🔥 Starting training (JSON AI)…\n")
-
-    ok = True
-    with open(LOG_FILE, "a") as log:
+    with open(LOG_FILE,"w") as log: log.write("🔥 Starting training (JSON AI)…\n")
+    ok=True
+    with open(LOG_FILE,"a") as log:
         if shards_folder:
             log.write(f"📂 Folder mode: {shards_folder}\n")
-            paths = sorted(glob.glob(os.path.join(shards_folder, "*.jsonl"))) + \
-                    sorted(glob.glob(os.path.join(shards_folder, "*.json"))) + \
-                    sorted(glob.glob(os.path.join(shards_folder, "*.jsonl.gz"))) + \
-                    sorted(glob.glob(os.path.join(shards_folder, "*.json.gz")))
+            paths = sorted(glob.glob(os.path.join(shards_folder,"*.jsonl"))) + \
+                    sorted(glob.glob(os.path.join(shards_folder,"*.json")))  + \
+                    sorted(glob.glob(os.path.join(shards_folder,"*.jsonl.gz"))) + \
+                    sorted(glob.glob(os.path.join(shards_folder,"*.json.gz")))
             if not paths:
-                log.write("❌ No shards found (*.jsonl / *.json / *.jsonl.gz / *.json.gz). Aborting.\n")
-                ok = False
+                log.write("❌ No shards found. Aborting.\n"); ok=False
             else:
-                tmp = "tmp_train.jsonl"
-                for i, p in enumerate(paths, 1):
-                    log.write(f"\n[{i}/{len(paths)}] Training on shard: {os.path.basename(p)}\n")
-                    if p.endswith(".gz"):
+                tmp="tmp_train.jsonl"
+                for i,pth in enumerate(paths,1):
+                    log.write(f"\n[{i}/{len(paths)}] Training on shard: {os.path.basename(pth)}\n")
+                    if pth.endswith(".gz"):
                         try:
-                            with gzip.open(p, "rt", encoding="utf-8") as rf, open(tmp, "w", encoding="utf-8") as wf:
-                                for line in rf:
-                                    wf.write(line)
-                            shard_path = tmp
+                            with gzip.open(pth,"rt",encoding="utf-8") as rf, open(tmp,"w",encoding="utf-8") as wf:
+                                for line in rf: wf.write(line)
+                            shard=tmp
                         except Exception as e:
-                            log.write(f"❌ Failed to read gz shard: {e}\n")
-                            ok = False
-                            break
+                            log.write(f"❌ Failed to read gz shard: {e}\n"); ok=False; break
                     else:
-                        shard_path = p
-                    if not _train_single_file(shard_path, log):
-                        ok = False
-                        break
+                        shard=pth
+                    if not _train_single_file(shard, log):
+                        ok=False; break
                 if os.path.exists(tmp):
                     try: os.remove(tmp)
                     except: pass
         else:
             if not dataset_path or not os.path.exists(dataset_path):
-                log.write("❌ Please upload a valid dataset first.\n")
-                ok = False
+                log.write("❌ Please upload a valid dataset first.\n"); ok=False
             else:
-                ok = _train_single_file(dataset_path, log)
+                ok=_train_single_file(dataset_path, log)
 
         if ok and os.path.isdir(MODEL_DIR):
             try:
-                time.sleep(0.5)  # settle delay for FS
+                time.sleep(0.5)
                 _zip_folder_atomic(MODEL_DIR, ZIP_FILE, ZIP_TEMP)
-                sz = _human_size(os.path.getsize(ZIP_FILE))
+                sz=_human_size(os.path.getsize(ZIP_FILE))
                 log.write(f"\n✅ Model zipped → {ZIP_FILE} ({sz})\n")
             except Exception as e:
                 log.write(f"\n❌ Zipping failed: {e}\n")
@@ -140,12 +160,11 @@ def _train_worker(dataset_path: str, shards_folder: str):
 
 def start_training(dataset_path: str, shards_folder: str):
     ensure_clean_zip()
-    t = threading.Thread(target=_train_worker, args=(dataset_path, shards_folder), daemon=True)
-    t.start()
+    threading.Thread(target=_train_worker, args=(dataset_path, shards_folder), daemon=True).start()
     return "🚀 Training started in the background. Use the Refresh buttons to update."
 
 def read_logs_once():
-    return _read_file_safely(LOG_FILE, "Waiting for logs...")
+    return _read_file_safely(LOG_FILE,"Waiting for logs...")
 
 def check_download():
     if os.path.exists(ZIP_FILE):
@@ -153,16 +172,13 @@ def check_download():
     else:
         return gr.update(visible=False, value=None), "No trained model yet."
 
-# ---- test helpers ----
+# --------- Test ----------
 def upload_test_model_zip(zip_file):
-    """Upload a model ZIP and extract it to models/test_<uuid>/ for testing."""
-    if zip_file is None:
-        return "❌ No file uploaded.", ""
+    if zip_file is None: return "❌ No file uploaded.", ""
     extract_root = os.path.join("models", f"test_{uuid.uuid4().hex}")
     os.makedirs(extract_root, exist_ok=True)
     try:
-        with zipfile.ZipFile(zip_file.name, "r") as zf:
-            zf.extractall(extract_root)
+        with zipfile.ZipFile(zip_file.name,"r") as zf: zf.extractall(extract_root)
         return f"✅ Model ZIP extracted to: {extract_root}", extract_root
     except Exception as e:
         return f"❌ Failed to extract: {e}", ""
@@ -171,36 +187,54 @@ def clear_uploaded_model():
     return "Model cleared. Will use trained_model/ if available.", ""
 
 def generate_response(prompt, uploaded_model_path):
-    if not prompt or not prompt.strip():
-        return "Please enter a prompt."
+    if not prompt or not prompt.strip(): return "Please enter a prompt."
     try:
         if uploaded_model_path and os.path.isdir(uploaded_model_path):
-            model_path = uploaded_model_path
-            src = "(uploaded model)"
+            model_path, src = uploaded_model_path, "(uploaded model)"
         elif os.path.isdir(MODEL_DIR):
-            model_path = MODEL_DIR
-            src = "(trained_model/)"
+            model_path, src = MODEL_DIR, "(trained_model/)"
         else:
-            model_path = "distilgpt2"
-            src = "(fallback: distilgpt2)"
+            model_path, src = "distilgpt2", "(fallback: distilgpt2)"
         gen = pipeline("text-generation", model=model_path, tokenizer="distilgpt2")
         out = gen(prompt, max_length=256, do_sample=True, temperature=0.7, truncation=True)[0]["generated_text"]
         return f"{out}\n\n— using {src}"
     except Exception as e:
         return f"❌ Error: {e}"
 
-# ---- UI ----
-with gr.Blocks(title="JSON AI Trainer") as app:
-    gr.Markdown("## 🧩 JSON AI Trainer\nUpload a dataset (JSONL/JSON), train in background, download the model, and test JSON-focused prompts.")
+# --------- UI ----------
+with gr.Blocks(title="JSON AI Trainer (with Dataset Generator)") as app:
+    gr.Markdown("## 🧩 JSON AI Trainer\nGenerate a large JSON dataset, train (single file or folder of shards), download the model, and test.")
 
     dataset_state = gr.State(value="")
     shard_folder_state = gr.State(value="")
     test_model_state = gr.State(value="")
 
+    with gr.Tab("🧪 Generate Dataset"):
+        with gr.Row():
+            total_in = gr.Number(value=1_000_000, label="Total samples")
+            shard_in = gr.Number(value=10_000, label="Rows per shard")
+        with gr.Row():
+            out_dir_in = gr.Textbox(value="json_dataset_v1", label="Output folder")
+            prefix_in = gr.Textbox(value="json", label="File prefix")
+        with gr.Row():
+            gen_btn = gr.Button("🚀 Start Generation")
+            gen_refresh_btn = gr.Button("🔁 Refresh Logs")
+        gen_status = gr.Textbox(label="Generator Status", interactive=False)
+        gen_logs = gr.Textbox(label="Generator Logs", lines=16)
+        with gr.Row():
+            list_folder = gr.Textbox(value="json_dataset_v1", label="Preview shards in folder")
+            list_btn = gr.Button("👀 List Shards")
+            list_out = gr.Textbox(label="Shard Preview", lines=8)
+
+        gen_btn.click(fn=start_generation, inputs=[total_in, shard_in, out_dir_in, prefix_in], outputs=gen_status
+        ).then(fn=read_gen_logs, outputs=gen_logs)
+        gen_refresh_btn.click(fn=read_gen_logs, outputs=gen_logs)
+        list_btn.click(fn=list_shards, inputs=list_folder, outputs=list_out)
+
     with gr.Tab("🧠 Train"):
-        gr.Markdown("Upload a single JSONL/JSON *or* provide a folder of shards (.jsonl, .json, .jsonl.gz, .json.gz).")
+        gr.Markdown("Upload a single JSON/JSONL file *or* train on a folder of shards (.json, .jsonl, .jsonl.gz, .json.gz).")
         with gr.Row():
-            file_input = gr.File(label="Upload single dataset file", file_types=[".jsonl", ".json"])
+            file_input = gr.File(label="Upload single dataset file", file_types=[".json",".jsonl"])
             upload_btn = gr.Button("📤 Upload (single file)")
         with gr.Row():
             shards_folder = gr.Textbox(value="", label="Folder with shards (optional)")
@@ -219,15 +253,9 @@ with gr.Blocks(title="JSON AI Trainer") as app:
             download_btn = gr.DownloadButton(label="📥 Download Trained Model (.zip)", visible=False, value=None)
 
         upload_btn.click(fn=upload_file, inputs=file_input, outputs=[status_box, dataset_state])
-        use_folder_btn.click(
-            fn=lambda p: ("✅ Using folder for training." if p.strip() else "❌ Provide a valid folder path.", p.strip()),
-            inputs=shards_folder,
-            outputs=[status_box, shard_folder_state]
-        )
-        start_btn.click(
-            fn=start_training,
-            inputs=[dataset_state, shard_folder_state],
-            outputs=status_box
+        use_folder_btn.click(fn=lambda p: ("✅ Using folder for training." if p.strip() else "❌ Provide a valid folder path.", p.strip()),
+                             inputs=shards_folder, outputs=[status_box, shard_folder_state])
+        start_btn.click(fn=start_training, inputs=[dataset_state, shard_folder_state], outputs=status_box
         ).then(fn=read_logs_once, outputs=log_output
         ).then(fn=check_download, outputs=[download_btn, download_info])
 
@@ -235,17 +263,13 @@ with gr.Blocks(title="JSON AI Trainer") as app:
         refresh_dl_btn.click(fn=check_download, outputs=[download_btn, download_info])
 
     with gr.Tab("🚀 Test"):
-        gr.Markdown("Upload a model ZIP or use the just-trained model. Try prompts like “Fix this JSON”, “Generate JSON schema for …”")
+        gr.Markdown("Upload a model ZIP or use the just-trained model.")
         with gr.Row():
             test_zip = gr.File(label="Upload Model ZIP", file_types=[".zip"])
             load_test_btn = gr.Button("📦 Load Uploaded Model ZIP")
             clear_test_btn = gr.Button("🧹 Clear Uploaded Model")
         test_status = gr.Textbox(label="Test Model Status", interactive=False)
-
-        prompt_input = gr.Textbox(
-            label="Prompt",
-            placeholder='e.g., "Generate valid JSON for a product with id, name, price, tags (array of strings)"'
-        )
+        prompt_input = gr.Textbox(label="Prompt", placeholder='e.g., "Generate JSON Schema for an invoice" or "Fix this JSON: {\'a\':1,}"')
         test_btn = gr.Button("🔍 Generate")
         response_output = gr.Textbox(label="AI Response", lines=12)
 
@@ -253,18 +277,16 @@ with gr.Blocks(title="JSON AI Trainer") as app:
         clear_test_btn.click(fn=clear_uploaded_model, outputs=[test_status, test_model_state])
         test_btn.click(fn=generate_response, inputs=[prompt_input, test_model_state], outputs=response_output)
 
-# Optional auto-start via env vars
-AUTOSTART = os.getenv("AUTOSTART_TRAIN", "0") == "1"
-AUTOSTART_DATASET = os.getenv("AUTOSTART_DATASET", "").strip()
-AUTOSTART_SHARDS  = os.getenv("AUTOSTART_SHARDS", "").strip()
+# Optional: autostart on boot via Space variables
+AUTOSTART = os.getenv("AUTOSTART_TRAIN","0") == "1"
+AUTOSTART_DATASET = os.getenv("AUTOSTART_DATASET","").strip()
+AUTOSTART_SHARDS  = os.getenv("AUTOSTART_SHARDS","").strip()
 if AUTOSTART and not os.path.exists(".autostart.started"):
-    open(".autostart.started", "w").close()
+    open(".autostart.started","w").close()
     try:
-        _ = start_training(AUTOSTART_DATASET if AUTOSTART_DATASET else "",
-                           AUTOSTART_SHARDS if AUTOSTART_SHARDS else "")
+        _ = start_training(AUTOSTART_DATASET if AUTOSTART_DATASET else "", AUTOSTART_SHARDS if AUTOSTART_SHARDS else "")
         _ = read_logs_once()
     except Exception as e:
-        with open(LOG_FILE, "a") as log:
-            log.write(f"\n❌ Autostart failed: {e}\n")
+        with open(LOG_FILE,"a") as log: log.write(f"\n❌ Autostart failed: {e}\n")
 
 app.launch()