Update app.py

app.py

@@ -1,413 +1,85 @@
-# app.py
-import os, shutil, subprocess, zipfile, traceback, io
+# app.py (minimal callback sanity check)
+import os, time
 from pathlib import Path
-from datetime import datetime
 import gradio as gr
 
-# ----------------- Paths -----------------
 ROOT = Path(__file__).resolve().parent
-DATA = ROOT / "dataset.jsonl"
-LOG  = ROOT / "train.log"
+LOG  = ROOT / "callback_test.log"
 RUNS = ROOT / "runs"
 RUNS.mkdir(exist_ok=True)
 
-# ----------------- Logging -----------------
 def append_log(msg: str):
-    msg = (msg or "").rstrip("\n")
     try:
-        with open(LOG, "a", encoding="utf-8") as lf:
-            lf.write(msg + "\n")
+        with open(LOG, "a", encoding="utf-8") as f:
+            f.write(msg.rstrip() + "\n")
     except Exception:
         pass
 
-def read_logs():
-    return LOG.read_text(encoding="utf-8")[-20000:] if LOG.exists() else "⏳ Waiting…"
+def ping():
+    append_log("PING called")
+    return "✅ Backend alive (ping)"
 
-# ----------------- Workspace & Models -----------------
-def ls_workspace() -> str:
-    rows = []
-    for p in sorted(ROOT.iterdir(), key=lambda x: (x.is_file(), x.name.lower())):
-        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 stream():
+    append_log("STREAM start")
+    for i in range(5):
+        time.sleep(0.3)
+        yield f"tick {i}"
+    append_log("STREAM done")
+    yield "✅ Stream ok"
 
 def list_models():
-    out = []
-    for base in [ROOT, RUNS]:
-        if not base.exists():
-            continue
-        for p in base.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))
-    return sorted(set(out))
-
-def dropdown_update_safe(models, prefer=None):
-    val = prefer if (prefer and prefer in models) else (models[0] if models else None)
-    return gr.update(choices=models, value=val)
-
-# ----------------- Dataset Upload -----------------
-def upload_dataset(file):
-    append_log("📥 upload_dataset clicked")
-    if not file:
-        return "❌ No file selected.", ls_workspace()
-    if hasattr(file, "name") and os.path.isfile(file.name):
-        shutil.copy(file.name, DATA)
-        return f"✅ Uploaded → {DATA.name}", ls_workspace()
-    return "⚠ Unexpected item; please upload a .jsonl file.", ls_workspace()
-
-# ----------------- Training (Live Logs) -----------------
-def start_training_live(run_name):
-    append_log("🚀 start_training_live clicked")
-    # Quick guard: dataset must exist
-    if not DATA.exists():
-        msg = "❌ dataset.jsonl not found. Upload a JSONL dataset first."
-        append_log(msg)
-        yield (msg, gr.update(value=None, visible=False), ls_workspace(), read_logs(), dropdown_update_safe(list_models()))
-        return
-
-    run_id = (run_name or "").strip() or datetime.now().strftime("run_%Y%m%d_%H%M%S")
-    out_dir  = RUNS / run_id
-    zip_path = RUNS / f"{run_id}.zip"
-
-    # clean only this run
-    if out_dir.exists():
-        shutil.rmtree(out_dir, ignore_errors=True)
-    if zip_path.exists():
-        zip_path.unlink()
-
-    # init log
-    LOG.write_text(f"🔥 Training started…\nRun: {run_id}\n", encoding="utf-8")
-    append_log(f"Workspace:\n{ls_workspace()}")
-
-    cmd = [
-        "python", str(ROOT / "train.py"),
-        "--dataset", str(DATA),
-        "--output",  str(out_dir),
-        "--zip_path", str(zip_path),
-        "--model_name", "Salesforce/codegen-350M-multi",
-        "--epochs", "1",
-        "--batch_size", "2",
-        "--block_size", "256",
-        "--learning_rate", "5e-5",
-    ]
-    append_log("▶ " + " ".join(cmd))
-
-    # start subprocess with live stdout
-    try:
-        proc = subprocess.Popen(
-            cmd,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.STDOUT,
-            bufsize=1,
-            universal_newlines=True,
-            encoding="utf-8",
-            errors="replace",
-        )
-    except Exception as e:
-        err = "❌ Failed to start train.py: " + "".join(traceback.format_exception_only(type(e), e))
-        append_log(err)
-        yield (err, gr.update(value=None, visible=False), ls_workspace(), read_logs(), dropdown_update_safe(list_models()))
-        return
-
-    live_log = io.StringIO()
-    status_msg = f"🚀 Training run '{run_id}' in progress…"
-
-    # stream loop
-    while True:
-        line = proc.stdout.readline()
-        if line == "" and proc.poll() is not None:
-            break
-        if line:
-            append_log(line.rstrip("\n"))
-            live_log.write(line)
-            text = live_log.getvalue()[-20000:]
-            yield (
-                status_msg,
-                gr.update(value=None, visible=False),
-                ls_workspace(),
-                text,
-                dropdown_update_safe(list_models(), prefer=None),
-            )
-            if zip_path.exists():
-                yield (
-                    "📦 Model zip created during run.",
-                    gr.update(value=str(zip_path), visible=True),
-                    ls_workspace(),
-                    text,
-                    dropdown_update_safe(list_models(), prefer=None),
-                )
-
-    code = proc.wait()
-
-    models = list_models()
-    model_update = dropdown_update_safe(models, prefer=str(out_dir) if out_dir.exists() else None)
-    final_logs = read_logs()
-
-    if code == 0 and zip_path.exists():
-        info = f"✅ Training complete. Saved: {out_dir.name} | Zip: {zip_path.name}"
-        append_log(info)
-        yield (info, gr.update(value=str(zip_path), visible=True), ls_workspace(), final_logs, model_update)
-    else:
-        info = f"❌ Training failed (exit {code}). Check logs below."
-        append_log(info)
-        yield (info, gr.update(value=None, visible=False), ls_workspace(), final_logs, model_update)
-
-def refresh_download():
-    append_log("↻ refresh_download clicked")
-    zips = sorted(RUNS.glob("*.zip"), key=lambda p: p.stat().st_mtime, reverse=True)
-    latest = zips[0] if zips else None
-    models = list_models()
-    return (
-        gr.update(value=(str(latest) if latest else None), visible=bool(latest)),
-        ls_workspace(),
-        dropdown_update_safe(models)
-    )
-
-# ----------------- Import a Zip as Model Folder -----------------
-def import_zip(zfile):
-    append_log("📦 import_zip clicked")
-    if not zfile:
-        return "❌ No zip selected.", list_models()
-    dest = ROOT / "imported_model"
-    if dest.exists():
-        shutil.rmtree(dest, ignore_errors=True)
-    dest.mkdir(parents=True, exist_ok=True)
-    with zipfile.ZipFile(zfile.name, "r") as z:
-        z.extractall(dest)
-    return f"✅ Imported to {dest.name}", list_models()
-
-# ----------------- Generation (cached pipeline) -----------------
-_GEN_CACHE = {"path": None, "pipe": None}
-
-def get_generation_pipeline(model_path: str):
-    from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-    import torch
-
-    if _GEN_CACHE["path"] == model_path and _GEN_CACHE["pipe"] is not None:
-        return _GEN_CACHE["pipe"]
-
-    append_log(f"🧩 Loading pipeline from: {model_path}")
-    tok = AutoTokenizer.from_pretrained(model_path, use_fast=True)
-    if tok.pad_token_id is None:
-        if tok.eos_token_id is not None:
-            tok.pad_token = tok.eos_token
-            append_log("ℹ No pad_token; using eos_token as pad_token.")
-        else:
-            tok.add_special_tokens({"pad_token": "[PAD]"})
-            append_log("ℹ Added [PAD] token to tokenizer.")
-    model = AutoModelForCausalLM.from_pretrained(model_path)
-    if getattr(model, "config", None) and getattr(model.config, "vocab_size", None) and len(tok) > model.config.vocab_size:
-        model.resize_token_embeddings(len(tok))
-        append_log(f"ℹ Resized embeddings to {len(tok)}.")
-
-    pipe = pipeline(
-        "text-generation",
-        model=model,
-        tokenizer=tok,
-        device_map="auto" if torch.cuda.is_available() else None,
-    )
-    _GEN_CACHE["path"] = model_path
-    _GEN_CACHE["pipe"] = pipe
-    append_log("✅ Pipeline loaded.")
-    return pipe
-
-# ----------------- Test Tab Helpers -----------------
-def ping():
-    append_log("🔔 Ping pressed (UI wiring OK)")
-    return "✅ UI is connected and responding."
-
-def load_selected_model(model_path):
-    append_log("📦 load_selected_model clicked")
-    # Dropdown may pass a list; coerce to string
-    if isinstance(model_path, list):
-        model_path = model_path[0] if model_path else None
-    if not model_path:
-        return "❌ Select a model first."
-    if not isinstance(model_path, str):
-        return f"❌ Invalid model path type: {type(model_path)._name_}"
-    p = Path(model_path)
+    # Just list subfolders under runs/ to prove dropdown gets values
+    items = [str(p) for p in sorted(RUNS.glob("*")) if p.is_dir()]
+    append_log(f"LIST_MODELS -> {items}")
+    return items
+
+def refresh_models():
+    return gr.update(choices=list_models(), value=None)
+
+def load_model(path):
+    append_log(f"LOAD_MODEL clicked: {path}")
+    if not path:
+        return "❌ Select a folder under runs/ (create one manually if empty)."
+    p = Path(path)
     if not p.exists() or not p.is_dir():
-        return f"❌ Model folder not found: {model_path}"
-    try:
-        append_log(f"📦 Load request → {model_path}")
-        _ = get_generation_pipeline(model_path)
-        append_log(f"✅ Loaded pipeline: {model_path}")
-        return f"✅ Loaded: {model_path}"
-    except Exception as e:
-        tb = traceback.format_exc()
-        append_log("❌ Load error:\n" + tb)
-        return "❌ Error while loading model:\n" + "".join(traceback.format_exception_only(type(e), e))
-
-def generate_once(model_path, prompt):
-    """Non-streaming fallback."""
-    append_log("▶ generate_once clicked")
-    # Coerce
-    if isinstance(model_path, list):
-        model_path = model_path[0] if model_path else None
+        return f"❌ Not found: {path}"
+    return f"✅ (Mock) loaded folder: {path}"
 
-    # validate
-    if not model_path:
-        msg = "❌ Select a model from the dropdown first."
-        append_log(msg); return msg
-    if not isinstance(model_path, str):
-        msg = f"❌ Invalid model path type: {type(model_path)._name_}"
-        append_log(msg); return msg
-    if not Path(model_path).exists():
-        msg = f"❌ Model folder not found: {model_path}"
-        append_log(msg); return msg
+def generate_once(path, prompt):
+    append_log(f"GENERATE clicked: path={path!r}, len(prompt)={len(prompt or '')}")
+    if not path:
+        return "❌ Pick a folder in the dropdown."
     if not prompt or not prompt.strip():
-        msg = "❌ Enter a prompt."
-        append_log(msg); return msg
+        return "❌ Enter a prompt."
+    # No model here—just echo to prove callback path works.
+    return f"🤖 MOCK RESPONSE\nModelFolder: {path}\nPrompt: {prompt.strip()[:120]}"
 
-    try:
-        pipe = get_generation_pipeline(model_path)
-        append_log(f"📝 Generating once… prompt_len={len(prompt)}")
-        result = pipe(
-            prompt.strip(),
-            max_new_tokens=80,
-            do_sample=True,
-            temperature=0.3,
-            top_p=0.9,
-            repetition_penalty=1.15,
-            no_repeat_ngram_size=4,
-            truncation=True,
-            return_full_text=True,
-        )
-        text = result[0].get("generated_text", "")
-        if not text:
-            append_log("⚠ Empty generated_text")
-            return "⚠ Model returned empty text. Try lowering temperature or adding more context."
-        append_log("✅ Generation OK.")
-        return text
-    except Exception as e:
-        tb = traceback.format_exc()
-        append_log("❌ Generation error:\n" + tb)
-        return "❌ Error during generation:\n" + "".join(traceback.format_exception_only(type(e), e))
+with gr.Blocks(title="Callback Sanity Check") as demo:
+    gr.Markdown("## 🔧 Minimal Callback Test\nIf these buttons do nothing, the issue is front-end/runtime (not your code).")
 
-def generate_stream(model_path, prompt):
-    """Streaming version (if Frontend streaming works)."""
-    yield "⏳ Loading model…"
-    append_log("▶ generate_stream clicked")
+    with gr.Row():
+        ping_btn = gr.Button("🔔 Ping")
+        stream_btn = gr.Button("📡 Stream Test")
+    out = gr.Textbox(label="Output", lines=8)
 
-    # Coerce
-    if isinstance(model_path, list):
-        model_path = model_path[0] if model_path else None
+    gr.Markdown("### Model list mock (folders under runs/)")
+    with gr.Row():
+        refresh_btn = gr.Button("↻ Refresh List")
+        model_dd = gr.Dropdown(choices=list_models(), label="Available Folders", interactive=True)
 
-    # validate
-    if not model_path:
-        msg = "❌ Select a model from the dropdown first."
-        append_log(msg); yield msg; return
-    if not isinstance(model_path, str):
-        msg = f"❌ Invalid model path type: {type(model_path)._name_}"
-        append_log(msg); yield msg; return
-    if not Path(model_path).exists():
-        msg = f"❌ Model folder not found: {model_path}"
-        append_log(msg); yield msg; return
-    if not prompt or not prompt.strip():
-        msg = "❌ Enter a prompt."
-        append_log(msg); yield msg; return
+    load_btn = gr.Button("📦 Load Folder (mock)")
+    load_status = gr.Textbox(label="Load Status", interactive=False)
 
-    try:
-        pipe = get_generation_pipeline(model_path)
-        yield "⚙ Generating… (this may take a bit on CPU)"
-        append_log(f"📝 Generating (stream)… prompt_len={len(prompt)}")
-        result = pipe(
-            prompt.strip(),
-            max_new_tokens=80,
-            do_sample=True,
-            temperature=0.3,
-            top_p=0.9,
-            repetition_penalty=1.15,
-            no_repeat_ngram_size=4,
-            truncation=True,
-            return_full_text=True,
-        )
-        text = result[0].get("generated_text", "")
-        if not text:
-            append_log("⚠ Empty generated_text")
-            yield "⚠ Model returned empty text. Try lowering temperature or adding more context."
-            return
-        append_log("✅ Generation OK.")
-        yield text
-    except Exception as e:
-        tb = traceback.format_exc()
-        append_log("❌ Generation error:\n" + tb)
-        yield "❌ Error during generation:\n" + "".join(traceback.format_exception_only(type(e), e))
+    prompt = gr.Textbox(label="Prompt", lines=4, placeholder="Type anything…")
+    gen_btn = gr.Button("Generate (mock)")
+    gen_out = gr.Textbox(label="Generated Text", lines=10)
 
-# ----------------- UI -----------------
-with gr.Blocks(title="Python AI — Train & Test") as app:
-    gr.Markdown("## 🧠 Python AI — Train & Test\n• Unique runs • Safe download • Cached generation • Live logs\n")
-
-    # ---------- Test Tab ----------
-    with gr.Tab("Test"):
-        gr.Markdown("### Choose a model folder or upload a .zip, then prompt it")
-        with gr.Row():
-            refresh_btn = gr.Button("↻ Refresh Model List")
-            ping_btn = gr.Button("🔔 Ping UI")  # sanity check
-        model_list = gr.Dropdown(
-            choices=list_models(),
-            label="Available AIs",
-            interactive=True,
-            allow_custom_value=True,
-            multiselect=False
-        )
-        load_btn = gr.Button("📦 Load Model")
-        load_status = gr.Textbox(label="Model Status", interactive=False)
-
-        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"
-        )
-        with gr.Row():
-            go_stream = gr.Button("Generate (stream)")
-            go_once   = gr.Button("Generate (once)")
-        out = gr.Textbox(label="AI Response", lines=20)
-
-    # ---------- Train Tab ----------
-    with gr.Tab("Train"):
-        with gr.Row():
-            ds = gr.File(label="📥 Upload JSONL", file_types=[".jsonl"])
-            ws = gr.Textbox(label="Workspace", lines=16, value=ls_workspace())
-        run_name = gr.Textbox(label="Run name (optional)", placeholder="e.g., python_small_v1")
-        up_status = gr.Textbox(label="Upload Status", interactive=False)
-        start = gr.Button("🚀 Start Training (Live Logs)", variant="primary")
-        logs = gr.Textbox(label="📜 Training Logs (live)", lines=18, value=read_logs())
-        status = gr.Textbox(label="Status", interactive=False)
-        download_file = gr.File(label="📦 Latest trained zip", visible=False)
-        refresh_dl_btn = gr.Button("Refresh Download")
-
-    # ---------- Wiring ----------
-    ds.change(upload_dataset, inputs=ds, outputs=[up_status, ws])
-
-    start.click(
-        start_training_live,
-        inputs=[run_name],
-        outputs=[status, download_file, ws, logs, model_list]
-    )
-
-    refresh_dl_btn.click(
-        refresh_download,
-        outputs=[download_file, ws, model_list]
-    )
-
-    refresh_btn.click(lambda: dropdown_update_safe(list_models()), outputs=model_list)
+    # Wiring
     ping_btn.click(ping, outputs=out)
-    load_btn.click(load_selected_model, inputs=[model_list], outputs=[load_status])
-    zip_in.change(import_zip, inputs=zip_in, outputs=[import_status, model_list])
-
-    # Generation (two modes)
-    go_stream.click(generate_stream, inputs=[model_list, prompt], outputs=out)
-    go_once.click(generate_once, inputs=[model_list, prompt], outputs=out)
+    stream_btn.click(stream, outputs=out)
+    refresh_btn.click(refresh_models, outputs=model_dd)
+    load_btn.click(load_model, inputs=model_dd, outputs=load_status)
+    gen_btn.click(generate_once, inputs=[model_dd, prompt], outputs=gen_out)
 
-# Critical: disable SSR; ensure queue is enabled
-app.queue(default_concurrency_limit=1)
-app.launch(ssr_mode=False, show_error=True)
+demo.queue(default_concurrency_limit=1)
+demo.launch(ssr_mode=False, show_error=True)