Create app.py

app.py

@@ -0,0 +1,171 @@
+import os, io, zipfile, shutil, subprocess, json, time, glob, tempfile
+import gradio as gr
+from pathlib import Path
+from typing import List, Tuple
+
+WORKDIR      = Path(".")
+DATASET_PATH = WORKDIR / "dataset.jsonl"
+LOG_PATH     = WORKDIR / "train.log"
+MODEL_DIR    = WORKDIR / "trained_model"        # training output folder
+ZIP_PATH     = WORKDIR / "trained_model.zip"    # zipped after train
+MODELS_ROOT  = WORKDIR                          # where we scan for saved AIs
+
+# ---------- helpers ----------
+def _safe_unzip(zip_file: str, out_dir: Path) -> str:
+    out_dir.mkdir(parents=True, exist_ok=True)
+    with zipfile.ZipFile(zip_file, "r") as z:
+        z.extractall(out_dir)
+    # return the inner model folder if zip contained a single directory
+    subdirs = [p for p in out_dir.iterdir() if p.is_dir()]
+    return str(subdirs[0] if len(subdirs) == 1 else out_dir)
+
+def _list_local_models() -> List[str]:
+    """
+    Return model folders found under MODELS_ROOT that look like HF models.
+    We include any folder that has a tokenizer.json or tokenizer_config.json.
+    """
+    candidates = []
+    for p in MODELS_ROOT.iterdir():
+        if not p.is_dir():
+            continue
+        if (p / "tokenizer.json").exists() or (p / "tokenizer_config.json").exists():
+            candidates.append(str(p))
+    return sorted(candidates)
+
+def _start_training_subprocess() -> int:
+    # clear old outputs
+    if MODEL_DIR.exists():
+        shutil.rmtree(MODEL_DIR)
+    if ZIP_PATH.exists():
+        ZIP_PATH.unlink(missing_ok=True)
+
+    cmd = [
+        "python", "train.py",
+        "--dataset", str(DATASET_PATH),
+        "--output",  str(MODEL_DIR),
+        # sensible defaults for quick, real training; adjust in train.py if needed
+        "--model_name", "Salesforce/codegen-350M-multi",
+        "--epochs", "1",
+        "--batch_size", "2",
+        "--block_size", "256",
+        "--learning_rate", "5e-5",
+    ]
+    LOG_PATH.write_text("🔥 Starting training...\n", encoding="utf-8")
+    with open(LOG_PATH, "a", encoding="utf-8") as lf:
+        proc = subprocess.Popen(cmd, stdout=lf, stderr=subprocess.STDOUT)
+        return proc.wait()
+
+def _zip_model_folder() -> bool:
+    if not MODEL_DIR.exists():
+        return False
+    if ZIP_PATH.exists():
+        ZIP_PATH.unlink()
+    shutil.make_archive(ZIP_PATH.with_suffix("").as_posix(), "zip", MODEL_DIR)
+    return ZIP_PATH.exists()
+
+# ---------- UI callbacks ----------
+def upload_dataset(file) -> str:
+    if file is None:
+        return "❌ No file selected."
+    shutil.copy(file.name, DATASET_PATH)
+    return f"✅ Uploaded: {file.name} → {DATASET_PATH.name}"
+
+def start_training() -> Tuple[str, str, gr.File]:
+    if not DATASET_PATH.exists():
+        return ("❌ Please upload a JSONL first.", "", gr.File.update(visible=False))
+    exit_code = _start_training_subprocess()
+
+    # after training, try to zip and expose
+    if exit_code == 0 and _zip_model_folder():
+        status = "✅ Training complete."
+        model_info = f"Saved: {MODEL_DIR.name}  |  Zip: {ZIP_PATH.name}"
+        return (status, model_info, gr.File.update(value=str(ZIP_PATH), visible=True))
+    else:
+        # surface the tail of the log for quick diagnosis
+        tail = ""
+        if LOG_PATH.exists():
+            with open(LOG_PATH, "r", encoding="utf-8") as f:
+                lines = f.readlines()[-30:]
+                tail = "".join(lines)
+        return (f"❌ Training failed (code {exit_code}).", tail, gr.File.update(visible=False))
+
+def read_logs() -> str:
+    if LOG_PATH.exists():
+        return LOG_PATH.read_text(encoding="utf-8")[-20_000:]  # last ~20k chars
+    return "⏳ Waiting for logs..."
+
+def refresh_model_list() -> List[str]:
+    return _list_local_models()
+
+def upload_model_zip(zip_file) -> Tuple[str, List[str]]:
+    if zip_file is None:
+        return "❌ No zip provided.", refresh_model_list()
+    out = WORKDIR / f"imported_{int(time.time())}"
+    path = _safe_unzip(zip_file.name, out)
+    msg = f"✅ Imported model at: {path}"
+    return msg, refresh_model_list()
+
+def generate(model_path: str, prompt: str) -> str:
+    if not model_path:
+        return "❌ Select a model."
+    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)
+
+        # decoding tuned for code
+        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"]
+        return out
+    except Exception as e:
+        return f"❌ Error: {e}"
+
+# ---------- UI ----------
+with gr.Blocks(title="Python AI Trainer") as demo:
+    gr.Markdown("## 🧠 Python AI Trainer\nUpload JSONL, train, then test your model.")
+
+    with gr.Tab("Train"):
+        file_in   = gr.File(label="📥 Upload JSONL Dataset", file_types=[".jsonl", ".jsonl.gz", ".json"])
+        up_status = gr.Textbox(label="Upload Status", interactive=False)
+        start_btn = gr.Button("🚀 Start Training", variant="primary")
+        logs_box  = gr.Textbox(label="📜 Live Logs (click Refresh)", lines=16)
+        refresh_logs = gr.Button("Refresh Logs")
+        status_box = gr.Textbox(label="Status", interactive=False)
+        model_info = gr.Textbox(label="Model Output", interactive=False)
+        dl = gr.File(label="📦 Download Trained Model (.zip)", visible=False)
+        refresh_dl = gr.Button("Refresh Download Area")
+
+        file_in.change(fn=upload_dataset, inputs=file_in, outputs=up_status)
+        start_btn.click(fn=start_training, outputs=[status_box, model_info, dl])
+        refresh_logs.click(fn=read_logs, outputs=logs_box)
+        refresh_dl.click(fn=lambda: (gr.File.update(value=str(ZIP_PATH), visible=ZIP_PATH.exists())),
+                         outputs=dl)
+
+    with gr.Tab("Test"):
+        gr.Markdown("### 🔬 Choose a stored AI and prompt it")
+        refresh_models_btn = gr.Button("↻ Refresh AI List")
+        model_list = gr.Dropdown(choices=_list_local_models(), label="Available AIs", interactive=True)
+        up_zip = gr.File(label="Or upload a model .zip to test", file_types=[".zip"])
+        zip_status = gr.Textbox(label="Model Import Status", interactive=False)
+        prompt = gr.Textbox(label="Prompt", lines=6, placeholder="### Instruction:\nPython: write a function ...\n### Response:\n")
+        generate_btn = gr.Button("Generate")
+        output = gr.Textbox(label="AI Response", lines=20)
+
+        refresh_models_btn.click(fn=refresh_model_list, outputs=model_list)
+        up_zip.change(fn=upload_model_zip, inputs=up_zip, outputs=[zip_status, model_list])
+        generate_btn.click(fn=generate, inputs=[model_list, prompt], outputs=output)
+
+demo.launch()