Spaces:
Sleeping
Sleeping
| import os | |
| import shutil | |
| import subprocess | |
| import threading | |
| import uuid | |
| import time | |
| import zipfile | |
| import glob | |
| import gzip | |
| import gradio as gr | |
| from transformers import pipeline | |
| # ---- Paths / constants ---- | |
| 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 corrupt downloads | |
| # ---- Helpers ---- | |
| 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 _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 _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 ensure_clean(): | |
| for p in (ZIP_FILE, ZIP_TEMP): | |
| if os.path.exists(p): | |
| try: | |
| os.remove(p) | |
| except: | |
| pass | |
| def _zip_folder_atomic(src_dir: str, zip_path: str, tmp_path: str): | |
| """Write to .part then rename → avoids corrupt/half-written zips.""" | |
| 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) | |
| # ============================================================ | |
| # DATASET GENERATOR (PYTHON) | |
| # ============================================================ | |
| def start_generation(total, shard_size, out_dir, prefix): | |
| """Kick off Python dataset generation in a background thread.""" | |
| total = int(total or 1_000_000) | |
| shard_size = int(shard_size or 10_000) | |
| out_dir = (out_dir or "python_dataset_v1").strip() | |
| prefix = (prefix or "python").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_python_dataset.py"): | |
| log.write("❌ make_python_dataset.py not found in repo root.\n") | |
| return | |
| try: | |
| proc = subprocess.Popen( | |
| [ | |
| "python", | |
| "make_python_dataset.py", | |
| "--total", str(total), | |
| "--shard_size", str(shard_size), | |
| "--out_dir", out_dir, | |
| "--prefix", prefix, | |
| ], | |
| stdout=log, | |
| stderr=subprocess.STDOUT, | |
| ) | |
| proc.wait() | |
| log.write(f"\n🔚 Generator exited with code {proc.returncode}\n") | |
| if proc.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...") | |
| def list_shards(folder): | |
| """Return a short preview of shard files (for sanity).""" | |
| if not folder or not os.path.isdir(folder): | |
| return "❌ Provide a valid folder path that contains .jsonl or .jsonl.gz shards." | |
| 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 (*.jsonl or *.jsonl.gz)." | |
| 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): | |
| """Copy uploaded dataset to a stable path; return status + saved 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): | |
| """Train once on a single JSON/JSONL 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...\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, "*.jsonl.gz"))) | |
| if not paths: | |
| 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 gz, stream to tmp jsonl | |
| 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 | |
| _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") | |
| return ok | |
| def start_training(dataset_path: str, shards_folder: str): | |
| ensure_clean() | |
| 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...") | |
| def check_download(): | |
| """Return download button state + info text (manual, non-streaming).""" | |
| 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 | |
| # ============================================================ | |
| def upload_test_model_zip(zip_file): | |
| """ | |
| Accept a model ZIP, extract to models/test_<uuid>/, return status + extracted path. | |
| ZIP should contain a HF model folder (config.json + tokenizer + weights). | |
| """ | |
| 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="Python AI Trainer (with Dataset Generator)") as app: | |
| gr.Markdown("## 🐍 Python AI Trainer\nGenerate a large Python dataset, train (single file or folder of shards), download the model, and test any model (uploaded or trained).") | |
| dataset_state = gr.State(value="") # path to single dataset file | |
| shard_folder_state = gr.State(value="") # folder containing shards | |
| test_model_state = gr.State(value="") | |
| # =============== Generate Dataset =============== | |
| with gr.Tab("🧪 Generate Dataset"): | |
| gr.Markdown("Generate a large Python dataset in shards (no streaming; use Refresh to see logs).") | |
| 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="python_dataset_v1", label="Output folder") | |
| prefix_in = gr.Textbox(value="python", 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="python_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) | |
| # ==================== Train ==================== | |
| with gr.Tab("🧠 Train"): | |
| gr.Markdown("Upload a single JSONL *or* provide a folder with shards (.jsonl / .jsonl.gz).") | |
| with gr.Row(): | |
| file_input = gr.File(label="Upload single JSONL dataset", file_types=[".jsonl"]) | |
| 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]) | |
| # ===================== Test ===================== | |
| with gr.Tab("🚀 Test"): | |
| gr.Markdown("Use an uploaded model ZIP or 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., Write a Python function that parses CSV and computes average") | |
| 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 on boot via env vars ---- | |
| AUTOSTART = os.getenv("AUTOSTART_TRAIN", "0") == "1" | |
| AUTOSTART_SINGLE_DATASET = os.getenv("AUTOSTART_DATASET", "").strip() | |
| AUTOSTART_SHARDS_FOLDER = os.getenv("AUTOSTART_SHARDS", "").strip() | |
| if AUTOSTART and not os.path.exists(".autostart.started"): | |
| open(".autostart.started", "w").close() | |
| try: | |
| _ = start_training(AUTOSTART_SINGLE_DATASET if AUTOSTART_SINGLE_DATASET else "", | |
| AUTOSTART_SHARDS_FOLDER if AUTOSTART_SHARDS_FOLDER 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() |