Create app.py

app.py

@@ -0,0 +1,270 @@
+import os
+import shutil
+import subprocess
+import threading
+import uuid
+import time
+import zipfile
+import gzip
+import glob
+import gradio as gr
+from transformers import pipeline
+
+# ---- constants / paths ----
+LOG_FILE = "train.log"
+MODEL_DIR = "trained_model"
+ZIP_FILE = "trained_model.zip"
+ZIP_TEMP = ZIP_FILE + ".part"  # atomic write to avoid half-written zips
+
+# ---- 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
+    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
+    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):
+            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)
+
+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)))
+    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
+
+# ---- training helpers ----
+def upload_file(file):
+    """Save uploaded dataset to uploads/ and return stable path."""
+    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
+
+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:
+        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")))
+            if not paths:
+                log.write("❌ No shards found (*.jsonl / *.json / *.jsonl.gz / *.json.gz). 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"):
+                        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
+                        except Exception as e:
+                            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
+                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
+            else:
+                ok = _train_single_file(dataset_path, log)
+
+        if ok and os.path.isdir(MODEL_DIR):
+            try:
+                time.sleep(0.5)  # settle delay for FS
+                _zip_folder_atomic(MODEL_DIR, ZIP_FILE, ZIP_TEMP)
+                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")
+        else:
+            log.write("\n❌ Training failed; no zip created.\n")
+
+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()
+    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...")
+
+def check_download():
+    if os.path.exists(ZIP_FILE):
+        return gr.update(visible=True, value=ZIP_FILE), _download_info_text()
+    else:
+        return gr.update(visible=False, value=None), "No trained model yet."
+
+# ---- test helpers ----
+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.", ""
+    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)
+        return f"✅ Model ZIP extracted to: {extract_root}", extract_root
+    except Exception as e:
+        return f"❌ Failed to extract: {e}", ""
+
+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."
+    try:
+        if uploaded_model_path and os.path.isdir(uploaded_model_path):
+            model_path = uploaded_model_path
+            src = "(uploaded model)"
+        elif os.path.isdir(MODEL_DIR):
+            model_path = MODEL_DIR
+            src = "(trained_model/)"
+        else:
+            model_path = "distilgpt2"
+            src = "(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.")
+
+    dataset_state = gr.State(value="")
+    shard_folder_state = gr.State(value="")
+    test_model_state = gr.State(value="")
+
+    with gr.Tab("🧠 Train"):
+        gr.Markdown("Upload a single JSONL/JSON *or* provide a folder of shards (.jsonl, .json, .jsonl.gz, .json.gz).")
+        with gr.Row():
+            file_input = gr.File(label="Upload single dataset file", file_types=[".jsonl", ".json"])
+            upload_btn = gr.Button("📤 Upload (single file)")
+        with gr.Row():
+            shards_folder = gr.Textbox(value="", label="Folder with shards (optional)")
+            use_folder_btn = gr.Button("📂 Use Folder For Training")
+        status_box = gr.Textbox(label="Status", interactive=False)
+
+        with gr.Row():
+            start_btn = gr.Button("🚀 Start Training")
+            refresh_btn = gr.Button("🔁 Refresh Logs")
+            refresh_dl_btn = gr.Button("📦 Refresh Download Area")
+        log_output = gr.Textbox(label="📜 Training Logs", lines=18)
+
+        with gr.Group():
+            gr.Markdown("### 📦 Trained Model")
+            download_info = gr.Markdown(value="No trained model yet.")
+            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
+        ).then(fn=read_logs_once, outputs=log_output
+        ).then(fn=check_download, outputs=[download_btn, download_info])
+
+        refresh_btn.click(fn=read_logs_once, outputs=log_output)
+        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 …”")
+        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)"'
+        )
+        test_btn = gr.Button("🔍 Generate")
+        response_output = gr.Textbox(label="AI Response", lines=12)
+
+        load_test_btn.click(fn=upload_test_model_zip, inputs=test_zip, outputs=[test_status, test_model_state])
+        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()
+if AUTOSTART and not os.path.exists(".autostart.started"):
+    open(".autostart.started", "w").close()
+    try:
+        _ = 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")
+
+app.launch()