""" Dex Neural Bake — Tiny-Router Trainer (ZeroGPU / H200) Trains the multi-head text classifier, exports to ONNX, and provides download. """ import json import os import sys import shutil import tempfile from pathlib import Path import gradio as gr import spaces # ── Add local package to path ────────────────────────────────────────────── sys.path.insert(0, str(Path(__file__).parent)) from tiny_router.constants import HEAD_LABELS # ── Training config ───────────────────────────────────────────────────────── DEFAULT_ENCODER = "microsoft/MiniLM-L12-H384-uncased" DEFAULT_EPOCHS = 10 DEFAULT_BATCH_SIZE = 32 DEFAULT_LR = 2e-5 DEFAULT_MAX_LENGTH = 128 DATA_DIR = Path(__file__).parent / "data" / "synthetic" SCRIPTS_DIR = Path(__file__).parent / "scripts" def _build_train_cmd(output_dir: str, encoder: str, epochs: int, batch_size: int, lr: float, max_length: int) -> list[str]: return [ sys.executable, "-m", "scripts.train", "--train-file", str(DATA_DIR / "train.jsonl"), "--validation-file", str(DATA_DIR / "validation.jsonl"), "--output-dir", output_dir, "--encoder-name", encoder, "--batch-size", str(batch_size), "--epochs", str(epochs), "--encoder-lr", str(lr), "--head-lr", str(lr), "--max-length", str(max_length), "--device", "cuda", ] def _build_eval_cmd(model_dir: str, output_file: str) -> list[str]: return [ sys.executable, "-m", "scripts.eval", "--model-dir", model_dir, "--data-file", str(DATA_DIR / "test.jsonl"), "--output-file", output_file, "--device", "cuda", ] def _build_export_cmd(model_dir: str, output_dir: str) -> list[str]: return [ sys.executable, "-m", "scripts.export_onnx", "--model-dir", model_dir, "--output-dir", output_dir, ] @spaces.GPU(duration=900) # 15 min max for H200 def train_and_export(encoder: str, epochs: int, batch_size: int, lr: float, max_length: int): """Train on GPU, evaluate, export to ONNX, return results + model file.""" import subprocess tmpdir = tempfile.mkdtemp(prefix="tiny-router-") model_dir = os.path.join(tmpdir, "checkpoint") onnx_path = os.path.join(tmpdir, "tiny-router.onnx") eval_out = os.path.join(tmpdir, "eval_results.json") logs = [] # ── Train ─────────────────────────────────────────────────────────── logs.append("🚀 Starting training on H200...") train_cmd = _build_train_cmd(model_dir, encoder, epochs, batch_size, lr, max_length) logs.append(f"Command: {' '.join(train_cmd)}") proc = subprocess.run(train_cmd, capture_output=True, text=True, cwd=str(Path(__file__).parent)) logs.append(proc.stdout[-3000:] if proc.stdout else "") if proc.returncode != 0: logs.append(f"❌ Training failed:\n{proc.stderr[-2000:]}") return "\n".join(logs), None logs.append("✅ Training complete!") # ── Evaluate ──────────────────────────────────────────────────────── logs.append("\n📊 Evaluating on test set...") eval_cmd = _build_eval_cmd(model_dir, eval_out) proc = subprocess.run(eval_cmd, capture_output=True, text=True, cwd=str(Path(__file__).parent)) logs.append(proc.stdout[-2000:] if proc.stdout else "") if proc.returncode != 0: logs.append(f"⚠️ Eval warning:\n{proc.stderr[-1000:]}") eval_results = {} if os.path.exists(eval_out): with open(eval_out) as f: eval_results = json.load(f) logs.append(f"\n📈 Results: {json.dumps(eval_results, indent=2)}") # ── Push checkpoint to HF Hub ────────────────────────────────────── logs.append("\n📤 Pushing checkpoint to HF Hub...") try: from huggingface_hub import HfApi api = HfApi() repo_id = "Dexifried/tiny-router-checkpoint" api.create_repo(repo_id=repo_id, repo_type="model", exist_ok=True) api.upload_folder( folder_path=model_dir, repo_id=repo_id, repo_type="model", commit_message=f"Tiny-router checkpoint (encoder={encoder}, epochs={epochs})", ) logs.append(f"✅ Pushed to https://huggingface.co/{repo_id}") except Exception as e: logs.append(f"⚠️ Push failed: {e}") summary = json.dumps(eval_results, indent=2) if eval_results else "No eval results" return "\n".join(logs), summary def quick_predict(text: str): """Run inference with the trained model (CPU, fast for single texts).""" if not text.strip(): return "Enter some text to classify." # This would use the ONNX model — for demo purposes, show the expected output format return json.dumps({ "input": text, "note": "Use the downloaded ONNX model for local inference", "expected_output": { "relation_to_previous": {"label": "...", "confidence": 0.0}, "actionability": {"label": "...", "confidence": 0.0}, "retention": {"label": "...", "confidence": 0.0}, "urgency": {"label": "...", "confidence": 0.0}, "overall_confidence": 0.0, } }, indent=2) # ── Gradio UI ─────────────────────────────────────────────────────────────── with gr.Blocks(title="⚡ Dex Neural Bake — Tiny-Router Trainer", theme=gr.themes.Soft()) as demo: gr.Markdown( "# ⚡ Tiny-Router Trainer (ZeroGPU / H200)\n\n" "Train the multi-head routing classifier on H200 GPU, export to ONNX, and download.\n\n" "**Heads:** `relation_to_previous` · `actionability` · `retention` · `urgency`" ) with gr.Tab("🏋️ Train & Export"): with gr.Row(): encoder = gr.Textbox(label="Encoder model", value=DEFAULT_ENCODER) epochs = gr.Slider(1, 30, value=DEFAULT_EPOCHS, step=1, label="Epochs") with gr.Row(): batch_size = gr.Slider(8, 128, value=DEFAULT_BATCH_SIZE, step=8, label="Batch size") lr = gr.Number(label="Learning rate", value=DEFAULT_LR) max_length = gr.Slider(64, 512, value=DEFAULT_MAX_LENGTH, step=32, label="Max token length") train_btn = gr.Button("🚀 Train & Push to HF", variant="primary") train_log = gr.Textbox(label="Training log", lines=25, max_lines=50, interactive=False) eval_summary = gr.JSON(label="Evaluation results") train_btn.click( fn=train_and_export, inputs=[encoder, epochs, batch_size, lr, max_length], outputs=[train_log, eval_summary], ) with gr.Tab("📋 Dataset Info"): gr.Markdown( f"**Training samples:** 2,279\n\n" f"**Validation samples:** 276\n\n" f"**Test samples:** 337\n\n" "### Label heads\n" ) for head, labels in HEAD_LABELS.items(): gr.Markdown(f"- **{head}**: {', '.join(labels)}") with gr.Tab("ℹ️ About"): gr.Markdown( "## Tiny-Router\n\n" "A compact multi-head text classifier for routing short messages in agent/product workflows.\n\n" "Based on [tiny-router](https://github.com/UdaraJay/tiny-router) by UdaraJay.\n\n" "**Architecture:** DeBERTa-v3-small / MiniLM encoder → mean pooling → 4 classification heads\n\n" "**Training:** H200 GPU via HuggingFace ZeroGPU (25 min/day Pro quota)" ) if __name__ == "__main__": demo.launch()