app.py

"""
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()