| import os, shutil, subprocess, threading, uuid, time, zipfile, gzip, glob |
| import gradio as gr |
| from transformers import pipeline |
|
|
| LOG_FILE = "train.log" |
| MODEL_DIR = "trained_model" |
| ZIP_FILE = "trained_model.zip" |
| ZIP_PART = ZIP_FILE + ".part" |
|
|
| def _human(n): |
| u=["B","KB","MB","GB"]; i=0; x=float(n) |
| while x>=1024 and i<len(u)-1: x/=1024; i+=1 |
| return f"{x:.1f} {u[i]}" |
|
|
| def _read(path, fb="Waiting..."): |
| try: |
| with open(path,"r",encoding="utf-8",errors="ignore") as f: return f.read() |
| except: return fb |
|
|
| def _zip_dir_atomic(src, out_path, tmp_path): |
| if os.path.exists(tmp_path): os.remove(tmp_path) |
| with zipfile.ZipFile(tmp_path,"w",zipfile.ZIP_DEFLATED) as z: |
| for root,_,files in os.walk(src): |
| for fn in files: |
| fp = os.path.join(root, fn) |
| z.write(fp, arcname=os.path.relpath(fp, src)) |
| if os.path.exists(out_path): os.remove(out_path) |
| os.replace(tmp_path, out_path) |
|
|
| def upload_file(f): |
| if f is None: return "❌ No file.", "" |
| os.makedirs("uploads", exist_ok=True) |
| dst = os.path.join("uploads", f"dataset_{uuid.uuid4().hex}.jsonl") |
| shutil.copy(f.name, dst) |
| return f"✅ Uploaded → {dst}", dst |
|
|
| def _train_single(dataset, log, quick, steps, subset): |
| args = ["python","train.py","--dataset",dataset,"--output",MODEL_DIR] |
| if quick: |
| args += ["--quick","1","--max_steps",str(steps or 8),"--subset",str(subset or 32)] |
| p = subprocess.Popen(args, stdout=log, stderr=subprocess.STDOUT) |
| p.wait() |
| log.write(f"\n ↳ train.py exited {p.returncode} for {os.path.basename(dataset)}\n") |
| return p.returncode == 0 |
|
|
| def _worker(dataset_path, shards_folder, quick, steps, subset): |
| with open(LOG_FILE,"w") as log: log.write("🔥 Starting training (C# 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,"*.jsonl.gz"))) |
| paths = [p for p in paths if "manifest" not in os.path.basename(p).lower()] |
| if not paths: |
| log.write("❌ No shards (*.jsonl / *.jsonl.gz).\n"); ok=False |
| else: |
| tmp="tmp_train.jsonl" |
| for i,pth in enumerate(paths,1): |
| log.write(f"\n[{i}/{len(paths)}] Shard: {os.path.basename(pth)}\n") |
| if pth.endswith(".gz"): |
| try: |
| 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"❌ GZ read failed: {e}\n"); ok=False; break |
| else: |
| shard = pth |
| if not _train_single(shard, log, quick, steps, subset): 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("❌ Upload a valid dataset.\n"); ok=False |
| else: |
| ok = _train_single(dataset_path, log, quick, steps, subset) |
|
|
| if ok and os.path.isdir(MODEL_DIR): |
| try: |
| _zip_dir_atomic(MODEL_DIR, ZIP_FILE, ZIP_PART) |
| sz = _human(os.path.getsize(ZIP_FILE)) |
| log.write(f"\n✅ Model zipped → {ZIP_FILE} ({sz})\n") |
| except Exception as e: |
| log.write(f"\n❌ Zip failed: {e}\n") |
| else: |
| log.write("\n❌ Training failed; no zip.\n") |
|
|
| def start_training(dataset_path, shards_folder, quick, steps, subset): |
| for p in (ZIP_FILE, ZIP_PART): |
| if os.path.exists(p): |
| try: os.remove(p) |
| except: pass |
| threading.Thread(target=_worker, args=(dataset_path, shards_folder, quick, steps, subset), daemon=True).start() |
| return "🚀 Training started. Use Refresh buttons." |
|
|
| def read_logs(): return _read(LOG_FILE, "Waiting for logs...") |
|
|
| def refresh_download(): |
| if os.path.exists(ZIP_FILE): |
| size=_human(os.path.getsize(ZIP_FILE)) |
| return gr.update(visible=True, value=ZIP_FILE), f"*Ready:* {ZIP_FILE} • {size}" |
| return gr.update(visible=False, value=None), "No trained model yet." |
|
|
| def load_test_zip(z): |
| if z is None: return "❌ No file.", "" |
| import zipfile |
| root = os.path.join("models", f"test_{uuid.uuid4().hex}") |
| os.makedirs(root, exist_ok=True) |
| try: |
| with zipfile.ZipFile(z.name,"r") as zz: zz.extractall(root) |
| return f"✅ Extracted to {root}", root |
| except Exception as e: |
| return f"❌ Extract failed: {e}", "" |
|
|
| def clear_test_model(): return "Cleared. Will use trained_model/ if present.", "" |
|
|
| def generate(prompt, model_path): |
| if not prompt.strip(): return "Enter a prompt." |
| try: |
| if model_path and os.path.isdir(model_path): m, src = model_path, "(uploaded)" |
| elif os.path.isdir(MODEL_DIR): m, src = MODEL_DIR, "(trained_model/)" |
| else: m, src = "sshleifer/tiny-gpt2", "(tiny fallback)" |
| gen = pipeline("text-generation", model=m, tokenizer="sshleifer/tiny-gpt2") |
| out = gen(prompt, max_length=120, 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}" |
|
|
| with gr.Blocks(title="C# AI Trainer (Quick Test Mode)") as app: |
| gr.Markdown("## ⚡ C# AI Trainer — includes *Quick Test* (few steps on a tiny model).") |
|
|
| ds_state = gr.State("") |
| folder_state = gr.State("") |
| test_model_state = gr.State("") |
|
|
| with gr.Tab("🧠 Train"): |
| with gr.Row(): |
| file_in = gr.File(label="Upload dataset (.jsonl)", file_types=[".jsonl"]) |
| up_btn = gr.Button("📤 Upload") |
| with gr.Row(): |
| shards_folder = gr.Textbox(value="", label="Folder with shards (optional)") |
| use_folder = gr.Button("📂 Use Folder") |
| with gr.Row(): |
| quick = gr.Checkbox(value=True, label="⚡ Quick Test (tiny model, few steps)") |
| steps = gr.Number(value=8, label="Max steps (Quick mode)", precision=0) |
| subset = gr.Number(value=32, label="Subset rows (Quick mode)", precision=0) |
| status = gr.Textbox(label="Status", interactive=False) |
| with gr.Row(): |
| start_btn = gr.Button("🚀 Start Training") |
| refresh_logs = gr.Button("🔁 Refresh Logs") |
| refresh_dl = gr.Button("📦 Refresh Download Area") |
| logs = gr.Textbox(label="📜 Logs", lines=18) |
| dl_btn = gr.DownloadButton(label="📥 Download Trained Model (.zip)", visible=False, value=None) |
| dl_info = gr.Markdown("No trained model yet.") |
|
|
| up_btn.click(fn=upload_file, inputs=file_in, outputs=[status, ds_state]) |
| use_folder.click(fn=lambda p: ("✅ Using folder." if p.strip() else "❌ Provide folder path.", p.strip()), |
| inputs=shards_folder, outputs=[status, folder_state]) |
| start_btn.click(fn=start_training, inputs=[ds_state, folder_state, quick, steps, subset], outputs=status |
| ).then(fn=read_logs, outputs=logs |
| ).then(fn=refresh_download, outputs=[dl_btn, dl_info]) |
| refresh_logs.click(fn=read_logs, outputs=logs) |
| refresh_dl.click(fn=refresh_download, outputs=[dl_btn, dl_info]) |
|
|
| with gr.Tab("🚀 Test"): |
| with gr.Row(): |
| zip_in = gr.File(label="Upload model ZIP", file_types=[".zip"]) |
| load_btn = gr.Button("📦 Load ZIP") |
| clear_btn = gr.Button("🧹 Clear") |
| test_status = gr.Textbox(label="Test Model Status", interactive=False) |
| prompt = gr.Textbox(label="Prompt", placeholder="e.g., Write a C# method that reverses a string.") |
| go = gr.Button("🔍 Generate") |
| out = gr.Textbox(label="AI Response", lines=12) |
|
|
| load_btn.click(fn=load_test_zip, inputs=zip_in, outputs=[test_status, test_model_state]) |
| clear_btn.click(fn=clear_test_model, outputs=[test_status, test_model_state]) |
| go.click(fn=generate, inputs=[prompt, test_model_state], outputs=out) |
|
|
| app.launch() |
