Update app.py

app.py

@@ -19,7 +19,7 @@ def _read(path, fb="Waiting..."):
 
 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:
+    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)
@@ -34,16 +34,16 @@ def upload_file(f):
     shutil.copy(f.name, dst)
     return f"✅ Uploaded → {dst}", dst
 
-def _train_single(dataset, log):
-    p = subprocess.Popen(
-        ["python","train.py","--dataset",dataset,"--output",MODEL_DIR],
-        stdout=log, stderr=subprocess.STDOUT
-    )
+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):
+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:
@@ -67,7 +67,7 @@ def _worker(dataset_path, shards_folder):
                             log.write(f"❌ GZ read failed: {e}\n"); ok=False; break
                     else:
                         shard = pth
-                    if not _train_single(shard, log): ok=False; break
+                    if not _train_single(shard, log, quick, steps, subset): ok=False; break
                 if os.path.exists(tmp):
                     try: os.remove(tmp)
                     except: pass
@@ -75,7 +75,7 @@ def _worker(dataset_path, shards_folder):
             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)
+                ok = _train_single(dataset_path, log, quick, steps, subset)
 
         if ok and os.path.isdir(MODEL_DIR):
             try:
@@ -87,16 +87,15 @@ def _worker(dataset_path, shards_folder):
         else:
             log.write("\n❌ Training failed; no zip.\n")
 
-def start_training(dataset_path, shards_folder):
-    try:
-        if os.path.exists(ZIP_FILE): os.remove(ZIP_FILE)
-        if os.path.exists(ZIP_PART): os.remove(ZIP_PART)
-    except: pass
-    threading.Thread(target=_worker, args=(dataset_path, shards_folder), daemon=True).start()
+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 read_logs(): return _read(LOG_FILE, "Waiting for logs...")
 
 def refresh_download():
     if os.path.exists(ZIP_FILE):
@@ -106,7 +105,7 @@ def refresh_download():
 
 def load_test_zip(z):
     if z is None: return "❌ No file.", ""
-    import zipfile, uuid
+    import zipfile
     root = os.path.join("models", f"test_{uuid.uuid4().hex}")
     os.makedirs(root, exist_ok=True)
     try:
@@ -115,29 +114,22 @@ def load_test_zip(z):
     except Exception as e:
         return f"❌ Extract failed: {e}", ""
 
-def clear_test_model():
-    return "Cleared. Will use trained_model/ if present.", ""
+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 = model_path
-            src="(uploaded)"
-        elif os.path.isdir(MODEL_DIR):
-            m = MODEL_DIR
-            src="(trained_model/)"
-        else:
-            m = "distilgpt2"           # tiny fallback
-            src="(fallback)"
-        gen = pipeline("text-generation", model=m, tokenizer="distilgpt2")
-        out = gen(prompt, max_length=200, do_sample=True, temperature=0.7, truncation=True)[0]["generated_text"]
+        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") as app:
-    gr.Markdown("## 🧩 C# AI Trainer — upload JSONL, train fast, download, and test.")
+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("")
@@ -150,6 +142,10 @@ with gr.Blocks(title="C# AI Trainer") as app:
         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")
@@ -162,7 +158,7 @@ with gr.Blocks(title="C# AI Trainer") as app:
         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], outputs=status
+        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)