app.py

#!/usr/bin/env python3
"""NanoCoder Interactive Web UI.

A simple Gradio-based frontend for:
- Training a tokenizer on a local codebase
- Preparing training data
- Launching / monitoring training
- Generating code completions from a trained model
- Fill-in-the-middle code infilling
"""

import os
import sys
import time
import json
import threading
import glob
from pathlib import Path

import gradio as gr
import torch

# Add project to path
sys.path.insert(0, os.path.dirname(__file__))

from nanocoder.config import ModelConfig, TrainConfig, PRESETS
from nanocoder.model import NanoCoder
from nanocoder.tokenizer import train_tokenizer, load_tokenizer
from nanocoder.data import prepare_memmap, PythonCodeDataset, MemmapDataset
from nanocoder.generate import load_model, complete, fill_in_middle

# ── Global state ──────────────────────────────────────────────

LOADED_MODEL = None
LOADED_TOKENIZER = None
LOADED_CONFIG = None
TRAIN_LOG = []
TRAINING_ACTIVE = False


# ── Tab 1: Setup & Data ──────────────────────────────────────

def scan_codebase(codebase_path):
    """Scan a directory and report Python file stats."""
    if not codebase_path or not os.path.isdir(codebase_path):
        return "Please enter a valid directory path."

    py_files = glob.glob(os.path.join(codebase_path, "**/*.py"), recursive=True)
    total_bytes = sum(os.path.getsize(f) for f in py_files if os.path.isfile(f))
    total_lines = 0
    for f in py_files[:500]:  # sample first 500
        try:
            with open(f, encoding="utf-8", errors="ignore") as fh:
                total_lines += sum(1 for _ in fh)
        except Exception:
            pass

    est_lines = total_lines * len(py_files) // max(min(len(py_files), 500), 1)
    est_tokens = total_bytes // 3  # rough: ~3 bytes per token for code

    report = f"""## Codebase Scan Results

| Metric | Value |
|--------|-------|
| Python files | {len(py_files):,} |
| Total size | {total_bytes / 1e6:.1f} MB |
| Estimated lines | ~{est_lines:,} |
| Estimated tokens | ~{est_tokens:,} |

{'This is a good size for training.' if est_tokens > 100_000 else 'Small codebase - consider supplementing with external data.'}

**Estimated training time (50M model):** ~{max(est_tokens / 2000 / 3600, 0.1):.1f} hours on your M4 Pro
"""
    return report


def train_tokenizer_ui(codebase_path, vocab_size):
    """Train a BPE tokenizer on the codebase."""
    if not codebase_path or not os.path.isdir(codebase_path):
        return "Enter a valid directory path first."

    output_dir = os.path.join(os.path.dirname(__file__), "output")
    os.makedirs(output_dir, exist_ok=True)
    tok_path = os.path.join(output_dir, "tokenizer.json")

    try:
        tok = train_tokenizer(
            data_paths=[codebase_path],
            vocab_size=int(vocab_size),
            output_path=tok_path,
        )
        # Test it
        sample = 'def hello():\n    print("Hello, world!")\n'
        encoded = tok.encode(sample)
        return (
            f"Tokenizer trained and saved to `{tok_path}`\n\n"
            f"Vocab size: **{tok.get_vocab_size():,}**\n\n"
            f"Sample encoding ({len(encoded.ids)} tokens):\n"
            f"```\n{encoded.tokens[:15]}\n```"
        )
    except Exception as e:
        return f"Error: {e}"


def prepare_data_ui(codebase_path):
    """Pre-tokenize the codebase into a memory-mapped binary."""
    output_dir = os.path.join(os.path.dirname(__file__), "output")
    tok_path = os.path.join(output_dir, "tokenizer.json")
    data_dir = os.path.join(output_dir, "data")
    os.makedirs(data_dir, exist_ok=True)
    train_bin = os.path.join(data_dir, "train.bin")

    if not os.path.exists(tok_path):
        return "Train a tokenizer first (Step 1)."
    if not codebase_path or not os.path.isdir(codebase_path):
        return "Enter a valid directory path."

    try:
        tok = load_tokenizer(tok_path)
        prepare_memmap([codebase_path], tok, train_bin)
        import numpy as np
        data = np.memmap(train_bin, dtype=np.uint16, mode="r")
        return (
            f"Data prepared and saved to `{train_bin}`\n\n"
            f"**Total tokens: {len(data):,}**\n"
            f"File size: {os.path.getsize(train_bin) / 1e6:.1f} MB\n"
            f"Training chunks (seq_len=1024): {len(data) // 1025:,}"
        )
    except Exception as e:
        return f"Error: {e}"


# ── Tab 2: Training ──────────────────────────────────────────

def start_training(preset, max_steps, learning_rate, batch_size, grad_accum, use_bf16):
    """Launch training in a background thread."""
    global TRAINING_ACTIVE, TRAIN_LOG

    if TRAINING_ACTIVE:
        return "Training is already running. Wait for it to finish."

    output_dir = os.path.join(os.path.dirname(__file__), "output")
    tok_path = os.path.join(output_dir, "tokenizer.json")
    data_dir = os.path.join(output_dir, "data")
    ckpt_dir = os.path.join(output_dir, "checkpoints")
    train_bin = os.path.join(data_dir, "train.bin")

    if not os.path.exists(tok_path):
        return "No tokenizer found. Complete the Setup tab first."
    if not os.path.exists(train_bin):
        return "No training data found. Complete the Setup tab first."

    tok = load_tokenizer(tok_path)
    model_config = PRESETS[preset]
    model_config.vocab_size = tok.get_vocab_size()

    train_config = TrainConfig(
        data_dir=data_dir,
        tokenizer_path=tok_path,
        batch_size=int(batch_size),
        gradient_accumulation_steps=int(grad_accum),
        max_steps=int(max_steps),
        learning_rate=float(learning_rate),
        use_bf16=bool(use_bf16),
        output_dir=ckpt_dir,
        log_interval=10,
        save_interval=max(int(max_steps) // 5, 100),
    )

    TRAIN_LOG = []
    TRAINING_ACTIVE = True

    def _train():
        global TRAINING_ACTIVE
        try:
            _run_training_with_log(model_config, train_config)
        except Exception as e:
            TRAIN_LOG.append(f"\nERROR: {e}")
        finally:
            TRAINING_ACTIVE = False

    thread = threading.Thread(target=_train, daemon=True)
    thread.start()
    return f"Training started! Model: {preset} ({model_config.param_count_estimate/1e6:.1f}M params), {max_steps} steps\n\nClick 'Refresh Log' to see progress."


def _run_training_with_log(model_config, train_config):
    """Training loop that writes to TRAIN_LOG instead of stdout."""
    import math
    from contextlib import nullcontext

    device = torch.device("cpu")
    use_bf16 = train_config.use_bf16
    dtype = torch.bfloat16 if use_bf16 else torch.float32
    autocast_ctx = (
        torch.amp.autocast(device_type="cpu", dtype=torch.bfloat16)
        if use_bf16 else nullcontext()
    )

    TRAIN_LOG.append(f"Device: CPU | Dtype: {dtype}")

    # Load data
    train_dataset = MemmapDataset(
        os.path.join(train_config.data_dir, "train.bin"),
        model_config.max_seq_len,
    )
    from nanocoder.data import get_dataloader
    train_loader = get_dataloader(train_dataset, batch_size=train_config.batch_size, shuffle=True)

    # Create model
    model = NanoCoder(model_config).to(device)
    params = model.param_count()
    TRAIN_LOG.append(f"Model: {params:,} parameters ({params/1e6:.1f}M)")

    # Optimizer
    decay_params = [p for n, p in model.named_parameters() if p.requires_grad and p.dim() >= 2]
    no_decay_params = [p for n, p in model.named_parameters() if p.requires_grad and p.dim() < 2]
    optimizer = torch.optim.AdamW([
        {"params": decay_params, "weight_decay": train_config.weight_decay},
        {"params": no_decay_params, "weight_decay": 0.0},
    ], lr=train_config.learning_rate, betas=(train_config.beta1, train_config.beta2))

    os.makedirs(train_config.output_dir, exist_ok=True)
    model_config.save(os.path.join(train_config.output_dir, "model_config.json"))

    model.train()
    data_iter = iter(train_loader)
    running_loss = 0.0
    t0 = time.time()
    tokens_processed = 0

    from nanocoder.train import get_lr

    for step in range(train_config.max_steps):
        if not TRAINING_ACTIVE:
            TRAIN_LOG.append("Training stopped by user.")
            break

        lr = get_lr(step, train_config)
        for pg in optimizer.param_groups:
            pg["lr"] = lr

        optimizer.zero_grad()
        accum_loss = 0.0

        for _ in range(train_config.gradient_accumulation_steps):
            try:
                batch = next(data_iter)
            except StopIteration:
                data_iter = iter(train_loader)
                batch = next(data_iter)

            input_ids = batch["input_ids"].to(device)
            targets = batch["targets"].to(device)

            with autocast_ctx:
                _, loss = model(input_ids, targets)
                loss = loss / train_config.gradient_accumulation_steps

            loss.backward()
            accum_loss += loss.item()
            tokens_processed += input_ids.numel()

        torch.nn.utils.clip_grad_norm_(model.parameters(), train_config.max_grad_norm)
        optimizer.step()
        running_loss += accum_loss

        if (step + 1) % train_config.log_interval == 0:
            elapsed = time.time() - t0
            tok_s = tokens_processed / elapsed if elapsed > 0 else 0
            avg = running_loss / train_config.log_interval
            ppl = math.exp(min(avg, 20))
            TRAIN_LOG.append(
                f"step {step+1:>6d}/{train_config.max_steps} | "
                f"loss {avg:.4f} | ppl {ppl:.1f} | "
                f"lr {lr:.2e} | {tok_s:.0f} tok/s"
            )
            running_loss = 0.0
            tokens_processed = 0
            t0 = time.time()

        if (step + 1) % train_config.save_interval == 0:
            path = os.path.join(train_config.output_dir, f"step_{step+1}.pt")
            torch.save({"model": model.state_dict(), "optimizer": optimizer.state_dict(),
                         "step": step + 1, "best_val_loss": float("inf")}, path)
            TRAIN_LOG.append(f"  >> checkpoint saved: {path}")

    # Final save
    final_path = os.path.join(train_config.output_dir, "final.pt")
    torch.save({"model": model.state_dict(), "optimizer": optimizer.state_dict(),
                 "step": train_config.max_steps, "best_val_loss": float("inf")}, final_path)
    TRAIN_LOG.append(f"\nTraining complete! Final checkpoint: {final_path}")


def get_train_log():
    """Return current training log."""
    if not TRAIN_LOG:
        return "No training log yet. Start training first."
    status = "RUNNING..." if TRAINING_ACTIVE else "FINISHED"
    return f"**Status: {status}**\n\n```\n" + "\n".join(TRAIN_LOG[-50:]) + "\n```"


def stop_training():
    global TRAINING_ACTIVE
    TRAINING_ACTIVE = False
    return "Stop signal sent. Training will halt after current step."


# ── Tab 3: Generate ──────────────────────────────────────────

def load_model_ui(checkpoint_name):
    """Load a trained model checkpoint."""
    global LOADED_MODEL, LOADED_TOKENIZER, LOADED_CONFIG

    output_dir = os.path.join(os.path.dirname(__file__), "output")
    ckpt_dir = os.path.join(output_dir, "checkpoints")
    tok_path = os.path.join(output_dir, "tokenizer.json")

    ckpt_path = os.path.join(ckpt_dir, checkpoint_name)
    if not os.path.exists(ckpt_path):
        return f"Checkpoint not found: {ckpt_path}"
    if not os.path.exists(tok_path):
        return "No tokenizer found."

    try:
        LOADED_TOKENIZER = load_tokenizer(tok_path)
        LOADED_MODEL, LOADED_CONFIG = load_model(ckpt_path)
        params = LOADED_MODEL.param_count()
        return f"Model loaded: **{params:,} parameters** ({params/1e6:.1f}M)\nCheckpoint: `{checkpoint_name}`"
    except Exception as e:
        return f"Error loading model: {e}"


def list_checkpoints():
    """List available checkpoints."""
    ckpt_dir = os.path.join(os.path.dirname(__file__), "output", "checkpoints")
    if not os.path.isdir(ckpt_dir):
        return gr.update(choices=["No checkpoints found"])
    files = sorted([f for f in os.listdir(ckpt_dir) if f.endswith(".pt")])
    if not files:
        return gr.update(choices=["No checkpoints found"])
    return gr.update(choices=files, value=files[-1])


def generate_code(prompt, max_tokens, temperature, top_k, top_p):
    """Generate code from a prompt."""
    if LOADED_MODEL is None or LOADED_TOKENIZER is None:
        return "Load a model first (click 'Load Model')."

    try:
        output, n_tokens, tok_per_sec = complete(
            LOADED_MODEL, LOADED_TOKENIZER, prompt,
            max_new_tokens=int(max_tokens),
            temperature=float(temperature),
            top_k=int(top_k),
            top_p=float(top_p),
        )
        stats = f"\n\n---\n*Generated {n_tokens} tokens at {tok_per_sec:.1f} tok/s*"
        return output + stats
    except Exception as e:
        return f"Error: {e}"


def infill_code(prefix, suffix, max_tokens, temperature):
    """Fill in code between prefix and suffix."""
    if LOADED_MODEL is None or LOADED_TOKENIZER is None:
        return "Load a model first."

    try:
        result = fill_in_middle(
            LOADED_MODEL, LOADED_TOKENIZER, prefix, suffix,
            max_new_tokens=int(max_tokens),
            temperature=float(temperature),
        )
        return result
    except Exception as e:
        return f"Error: {e}"


# ── Build the UI ──────────────────────────────────────────────

def build_app():
    with gr.Blocks(title="NanoCoder") as app:
        gr.Markdown(
            "# NanoCoder\n"
            "Train a tiny Python coding model on your CPU, then use it to generate code."
        )

        with gr.Tabs():
            # ── Tab 1: Setup ──
            with gr.Tab("1. Setup & Data"):
                gr.Markdown("### Point NanoCoder at a Python codebase to train on")
                with gr.Row():
                    codebase_input = gr.Textbox(
                        label="Codebase Path",
                        placeholder="/path/to/your/python/project",
                        scale=3,
                    )
                    scan_btn = gr.Button("Scan", scale=1)
                scan_output = gr.Markdown()
                scan_btn.click(scan_codebase, inputs=codebase_input, outputs=scan_output)

                gr.Markdown("---\n### Step 1: Train Tokenizer")
                with gr.Row():
                    vocab_size_input = gr.Slider(
                        4096, 32000, value=16384, step=1024,
                        label="Vocabulary Size",
                    )
                    tok_btn = gr.Button("Train Tokenizer")
                tok_output = gr.Markdown()
                tok_btn.click(train_tokenizer_ui, inputs=[codebase_input, vocab_size_input], outputs=tok_output)

                gr.Markdown("---\n### Step 2: Prepare Training Data")
                data_btn = gr.Button("Prepare Data")
                data_output = gr.Markdown()
                data_btn.click(prepare_data_ui, inputs=codebase_input, outputs=data_output)

            # ── Tab 2: Train ──
            with gr.Tab("2. Train"):
                gr.Markdown("### Configure and launch training")
                with gr.Row():
                    preset_input = gr.Radio(
                        ["50m", "100m", "200m"], value="50m",
                        label="Model Size",
                    )
                    bf16_input = gr.Checkbox(value=False, label="BFloat16 (OFF recommended - fp32 is faster on CPU)")

                with gr.Row():
                    steps_input = gr.Slider(100, 100000, value=5000, step=100, label="Max Steps")
                    lr_input = gr.Number(value=3e-4, label="Learning Rate")

                with gr.Row():
                    bs_input = gr.Slider(1, 8, value=1, step=1, label="Micro Batch Size")
                    ga_input = gr.Slider(1, 64, value=16, step=1, label="Gradient Accumulation Steps")

                with gr.Row():
                    train_btn = gr.Button("Start Training", variant="primary")
                    stop_btn = gr.Button("Stop Training", variant="stop")
                    refresh_btn = gr.Button("Refresh Log")

                train_output = gr.Markdown()
                log_output = gr.Markdown()

                train_btn.click(
                    start_training,
                    inputs=[preset_input, steps_input, lr_input, bs_input, ga_input, bf16_input],
                    outputs=train_output,
                )
                stop_btn.click(stop_training, outputs=train_output)
                refresh_btn.click(get_train_log, outputs=log_output)

            # ── Tab 3: Generate ──
            with gr.Tab("3. Generate"):
                gr.Markdown("### Generate Python code with your trained model")

                with gr.Row():
                    ckpt_dropdown = gr.Dropdown(label="Checkpoint", choices=[], interactive=True)
                    refresh_ckpts_btn = gr.Button("Refresh")
                    load_btn = gr.Button("Load Model", variant="primary")
                load_output = gr.Markdown()

                refresh_ckpts_btn.click(list_checkpoints, outputs=ckpt_dropdown)
                load_btn.click(load_model_ui, inputs=ckpt_dropdown, outputs=load_output)

                gr.Markdown("---")

                with gr.Tabs():
                    with gr.Tab("Code Completion"):
                        prompt_input = gr.Code(
                            language="python",
                            label="Prompt",
                            value="def fibonacci(n):\n    ",
                            lines=10,
                        )
                        with gr.Row():
                            max_tok = gr.Slider(32, 512, value=256, step=32, label="Max Tokens")
                            temp = gr.Slider(0.1, 2.0, value=0.8, step=0.1, label="Temperature")
                            topk = gr.Slider(0, 100, value=50, step=5, label="Top-K")
                            topp = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-P")
                        gen_btn = gr.Button("Generate", variant="primary")
                        gen_output = gr.Code(language="python", label="Output", lines=15)
                        gen_btn.click(
                            generate_code,
                            inputs=[prompt_input, max_tok, temp, topk, topp],
                            outputs=gen_output,
                        )

                    with gr.Tab("Fill-in-the-Middle"):
                        gr.Markdown("Enter code before and after the cursor. The model fills the gap.")
                        fim_prefix = gr.Code(
                            language="python", label="Prefix (before cursor)",
                            value="def add(a, b):\n", lines=5,
                        )
                        fim_suffix = gr.Code(
                            language="python", label="Suffix (after cursor)",
                            value="\n\nresult = add(3, 4)", lines=5,
                        )
                        with gr.Row():
                            fim_max_tok = gr.Slider(16, 256, value=128, step=16, label="Max Tokens")
                            fim_temp = gr.Slider(0.1, 2.0, value=0.8, step=0.1, label="Temperature")
                        fim_btn = gr.Button("Infill", variant="primary")
                        fim_output = gr.Code(language="python", label="Generated Infill", lines=5)
                        fim_btn.click(
                            infill_code,
                            inputs=[fim_prefix, fim_suffix, fim_max_tok, fim_temp],
                            outputs=fim_output,
                        )

            # ── Tab 4: About ──
            with gr.Tab("About"):
                gr.Markdown("""
## NanoCoder

A tiny Python coding model you can train on your own CPU.

### Architecture
- Decoder-only transformer with RoPE, GQA, SwiGLU, RMSNorm
- Presets: 50M (25M actual), 100M (62M actual), 200M (176M actual)
- Fill-in-the-middle (FIM) support for code completion

### How It Works
1. **Setup**: Point it at a Python codebase, train a tokenizer, prepare data
2. **Train**: Train the model on your CPU (M4 Pro w/ bf16 = fast)
3. **Generate**: Use the trained model for code completion

### Inspired By
- **phi-1** (Microsoft): Data quality > quantity
- **SmolLM** (HuggingFace): Depth > width for small models
- **nanoGPT** (Karpathy): Clean, minimal implementation
""")

        # Auto-refresh checkpoints on load
        app.load(list_checkpoints, outputs=ckpt_dropdown)

    return app


if __name__ == "__main__":
    app = build_app()
    app.launch(
        server_name="127.0.0.1", server_port=7860, share=False,
        theme=gr.themes.Soft(primary_hue="blue"),
    )