Remove nested directory: BitTransformerLM/bit_transformer/dashboard_app.py

BitTransformerLM/bit_transformer/dashboard_app.py

@@ -1,927 +0,0 @@
-import io
-import json
-import os
-import traceback
-import inspect
-from typing import Any, Dict, List
-
-from flask import Flask, jsonify, request, render_template, send_file
-import subprocess
-import sys
-import warnings
-import matplotlib.pyplot as plt
-import torch
-import torch.nn.functional as F
-import requests
-import gzip
-
-from .model import BitTransformerLM, infer_long_sequence
-from .optimization import configure_optimizer
-from .collapse import collapse_submodel
-from .dashboard import plot_telemetry
-from .scale import expand_model
-from .bit_io import text_to_bits, bits_to_text
-from .safety import hil_safe_inference
-from .compression import model_output_decompress, compress_bits
-from .distributed import wrap_fsdp
-from .training import train_loop
-from .telemetry import detect_metric_drift
-from .quantization import prepare_qat_fx, convert_qat_fx
-from torch.distributed.fsdp import FullyShardedDataParallel
-from .hf_checkpoint import hf_login, save_checkpoint, download_checkpoint
-
-
-app = Flask(__name__)
-app.config["MAX_CONTENT_LENGTH"] = 1 * 1024 * 1024  # 1MB upload limit
-
-MCP_SERVER_ADDR = os.getenv("MCP_SERVER_ADDR")
-
-
-@app.errorhandler(Exception)
-def handle_exception(err):
-    """Return JSON error responses with stack traces."""
-    return (
-        jsonify({"error": str(err), "trace": traceback.format_exc()}),
-        getattr(err, "code", 500),
-    )
-
-class MetricDriftWarning(UserWarning):
-    """Raised when telemetry metrics drift beyond the configured threshold."""
-
-def _switch_torch(use_gpu: bool) -> None:
-    """Install the appropriate PyTorch wheel and restart the process."""
-    have_cuda = torch.version.cuda is not None
-    if use_gpu == have_cuda:
-        return
-    wheel = "torch==2.7.1+cu118" if use_gpu else "torch==2.7.1+cpu"
-    url = "https://download.pytorch.org/whl/cu118" if use_gpu else "https://download.pytorch.org/whl/cpu"
-    subprocess.run([
-        sys.executable,
-        "-m",
-        "pip",
-        "install",
-        "--extra-index-url",
-        url,
-        wheel,
-    ], check=True)
-    os.execv(sys.executable, [sys.executable] + sys.argv)
-
-def mcp_post(path: str, data=None):
-    if not MCP_SERVER_ADDR:
-        return None
-    url = MCP_SERVER_ADDR.rstrip("/") + path
-    resp = requests.post(url, json=data)
-    resp.raise_for_status()
-    if resp.headers.get("Content-Type", "").startswith("image/"):
-        return resp.content
-    return resp.json()
-
-def mcp_get(path: str):
-    if not MCP_SERVER_ADDR:
-        return None
-    url = MCP_SERVER_ADDR.rstrip("/") + path
-    resp = requests.get(url)
-    resp.raise_for_status()
-    if resp.headers.get("Content-Type", "").startswith("image/"):
-        return resp.content
-    return resp.json()
-
-class ModelManager:
-    """Manage model state and training utilities for the dashboard."""
-
-    def __init__(
-        self,
-        snapshot_dir: str | None = None,
-        telemetry_log: str | None = None,
-        *,
-        drift_window: int = 10,
-        drift_threshold: float = 0.2,
-    ) -> None:
-        self.snapshot_dir = snapshot_dir or os.getenv("SNAPSHOT_DIR", "snapshots")
-        self.telemetry_log = telemetry_log or os.getenv("TELEMETRY_LOG")
-        if self.telemetry_log is None:
-            self.telemetry_log = os.path.join(self.snapshot_dir, "metrics.json")
-        os.makedirs(self.snapshot_dir, exist_ok=True)
-        self.weights_path = os.path.join(self.snapshot_dir, "model.pt")
-
-        self.model: BitTransformerLM | None = None
-        self.optimizer: torch.optim.Optimizer | None = None
-        self.scheduler: torch.optim.lr_scheduler._LRScheduler | None = None
-        self.total_steps = 100
-        self.metrics: Dict[str, List[float]] = {
-            "negentropy_logits": [],
-            "lz_complexity_logits": [],
-            "symbiosis_score": [],
-        }
-        self.drift_window = drift_window
-        self.drift_threshold = drift_threshold
-        self.lambda_K = 1.0
-        self.lambda_C = 1.0
-        self.lambda_S = 1.0
-        self.c_floor = 0.3
-        self.s_floor = 0.5
-        self.causal = True
-        self.diffusion = False
-        self.decompress_output = False
-        self.use_compression = False
-        self.use_gpu = False
-        self.qat = False
-
-        # Load any existing state
-        if os.path.exists(self.telemetry_log):
-            try:
-                with open(self.telemetry_log) as f:
-                    saved = json.load(f)
-                for key in self.metrics:
-                    self.metrics[key] = saved.get(key, [])
-            except Exception:
-                pass
-        if os.path.exists(self.weights_path):
-            try:
-                self.model = torch.load(self.weights_path, map_location="cpu")
-                self.optimizer, self.scheduler = configure_optimizer(
-                    self.model, lr=1e-3, total_steps=self.total_steps
-                )
-                self._apply_device()
-            except Exception:
-                self.model = None
-
-        config_path = os.getenv("MODEL_CONFIG", "/config/model_params.json")
-        if self.model is None and os.path.exists(config_path):
-            try:
-                with open(config_path) as f:
-                    params = json.load(f)
-                self.init_model(params)
-            except Exception:
-                pass
-
-    def init_model(self, params: Dict) -> None:
-        int_fields = {
-            "d_model",
-            "nhead",
-            "num_layers",
-            "dim_feedforward",
-            "max_seq_len",
-            "chunk_size",
-            "overlap",
-        }
-        float_fields = {"act_threshold"}
-        bool_fields = {"reversible", "use_checkpoint"}
-        clean: Dict[str, Any] = {}
-        for k, v in params.items():
-            if v is None:
-                clean[k] = None
-            elif k in int_fields:
-                clean[k] = int(v)
-            elif k in float_fields:
-                clean[k] = float(v)
-            elif k in bool_fields:
-                clean[k] = bool(v)
-            else:
-                clean[k] = v
-        self.model = BitTransformerLM(
-            **clean,
-            lambda_K=self.lambda_K,
-            lambda_C=self.lambda_C,
-            lambda_S=self.lambda_S,
-        )
-        self.optimizer, self.scheduler = configure_optimizer(
-            self.model, lr=1e-3, total_steps=self.total_steps
-        )
-        self._apply_device()
-        for key in self.metrics:
-            self.metrics[key].clear()
-
-    def set_lambdas(self, k: float, c: float, s: float) -> None:
-        """Update λ weights and propagate to the model."""
-        self.lambda_K = k
-        self.lambda_C = c
-        self.lambda_S = s
-        if self.model is not None:
-            self.model.set_lambdas(k, c, s)
-
-    def set_floors(self, c_floor: float, s_floor: float) -> None:
-        """Update safety floors for complexity (C) and symbiosis (S)."""
-        self.c_floor = c_floor
-        self.s_floor = s_floor
-
-    def set_diffusion(self, flag: bool) -> None:
-        """Toggle Diffusion LM mode which disables causal masking and chunking."""
-        self.diffusion = flag
-        self.causal = not flag
-        if self.model is not None and flag:
-            self.model.chunk_size = None
-
-    def set_decompress_output(self, flag: bool) -> None:
-        """Enable or disable decompression of model outputs."""
-        self.decompress_output = flag
-
-    def set_compression(self, flag: bool) -> None:
-        """Toggle automatic compression of inputs."""
-        self.use_compression = flag
-
-    def set_qat(self, flag: bool) -> None:
-        """Enable or disable 4-bit quantization-aware training."""
-        self.qat = flag
-        if self.model is None:
-            return
-        if flag:
-            self.model = prepare_qat_fx(self.model)
-        else:
-            self.model = convert_qat_fx(self.model)
-
-    def set_gpu(self, flag: bool) -> None:
-        """Toggle GPU acceleration and FSDP, reinstalling PyTorch if needed."""
-        _switch_torch(flag)
-        self.use_gpu = flag and torch.cuda.is_available()
-        self._apply_device()
-
-    def _apply_device(self) -> None:
-        """Move the model to the selected device and wrap with FSDP if needed."""
-        if self.model is None:
-            return
-        if self.use_gpu:
-            device = torch.device("cuda")
-            if isinstance(self.model, FullyShardedDataParallel):
-                base = self.model.module
-            else:
-                base = self.model
-            base = base.to(device)
-            self.model = wrap_fsdp(base, device_id=device)
-        else:
-            device = torch.device("cpu")
-            if isinstance(self.model, FullyShardedDataParallel):
-                self.model = self.model.module
-            self.model = self.model.to(device)
-
-    def train_step(self, bits: torch.Tensor) -> float:
-        assert (
-            self.model is not None
-            and self.optimizer is not None
-            and self.scheduler is not None
-        )
-        self.model.train()
-        device = next(self.model.parameters()).device
-        bits = bits.to(device)
-        ratio = 1.0
-        if self.use_compression:
-            comps = [compress_bits(row.to(torch.uint8)) for row in bits]
-            comp_len = sum(c.numel() for c in comps)
-            ratio = min(comp_len / bits.numel(), 1.0)
-            logits, telemetry = self.model.forward_compressed(comps, causal=self.causal)
-        else:
-            logits, telemetry = self.model(bits, causal=self.causal)
-        pred = logits[:, :-1, :].reshape(-1, 2)
-        target = bits[:, 1:].reshape(-1)
-        loss = F.cross_entropy(pred, target)
-        loss.backward()
-        torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
-        self.optimizer.step()
-        self.scheduler.step()
-        self.optimizer.zero_grad()
-        self._log_metrics(telemetry)
-        self._save_state()
-        return loss.item(), ratio
-
-    def train_epochs(
-        self,
-        bits: torch.Tensor,
-        *,
-        epochs: int = 1,
-        compress_prob: float = 0.5,
-        direct_prob: float = 0.0,
-        batch_size: int = 8,
-        num_workers: int = 0,
-        accum_steps: int = 1,
-        amp: bool = False,
-        compile_model: bool = False,
-    ) -> List[Dict[str, float]]:
-        """Run ``train_loop`` on a batch tensor and persist the state."""
-        assert self.model is not None
-        device = next(self.model.parameters()).device
-        bits = bits.to(device)
-        import math
-        steps_per_epoch = max(1, math.ceil(len(bits) / batch_size))
-        self.total_steps = math.ceil(epochs * steps_per_epoch / accum_steps)
-        self.optimizer, self.scheduler = configure_optimizer(
-            self.model, lr=1e-3, total_steps=self.total_steps
-        )
-        metrics = train_loop(
-            self.model,
-            bits,
-            epochs=epochs,
-            compress_prob=compress_prob if self.use_compression else 0.0,
-            direct_prob=direct_prob,
-            batch_size=batch_size,
-            num_workers=num_workers,
-            accum_steps=accum_steps,
-            amp=amp,
-            compile_model=compile_model,
-            forward_kwargs={"causal": self.causal},
-            optimizer=self.optimizer,
-            scheduler=self.scheduler,
-        )
-        self._save_state()
-        return metrics
-
-    def scale_up(self, width_mult: float = 1.0) -> None:
-        assert self.model is not None
-        params = dict(
-            d_model=int(self.model.d_model * width_mult),
-            nhead=self.model.layers[0].self_attn.num_heads,
-            num_layers=self.model.num_layers * 2,
-            dim_feedforward=int(self.model.layers[0].linear1.out_features * width_mult),
-            max_seq_len=self.model.pos_enc.pe.size(0),
-        )
-        self.model = expand_model(self.model, {
-            **params,
-            "lambda_K": self.lambda_K,
-            "lambda_C": self.lambda_C,
-            "lambda_S": self.lambda_S,
-        })
-        self.optimizer, self.scheduler = configure_optimizer(
-            self.model, lr=1e-3, total_steps=self.total_steps
-        )
-        self._save_state()
-
-    def collapse(self, cluster_bits: List[List[int]], target_params: Dict, width_scale: float = 1.0) -> None:
-        self.model, _ = collapse_submodel(
-            cluster_bits,
-            target_params,
-            width_scale=width_scale,
-            forward_kwargs={"causal": self.causal},
-        )
-        self.model.set_lambdas(self.lambda_K, self.lambda_C, self.lambda_S)
-        self.optimizer, self.scheduler = configure_optimizer(
-            self.model, lr=1e-3, total_steps=self.total_steps
-        )
-        self._apply_device()
-        for key in self.metrics:
-            self.metrics[key].clear()
-
-    def infer(self, bits: torch.Tensor) -> Dict:
-        assert self.model is not None
-        self.model.eval()
-        device = next(self.model.parameters()).device
-        bits = bits.to(device)
-        ratio = 1.0
-        with torch.no_grad():
-            if self.use_compression:
-                comps = [compress_bits(row.to(torch.uint8)) for row in bits]
-                comp_len = sum(c.numel() for c in comps)
-                ratio = min(comp_len / bits.numel(), 1.0)
-                logits, telemetry = self.model.forward_compressed(comps, causal=self.causal)
-            else:
-                logits, telemetry = self.model(bits, causal=self.causal)
-        self._log_metrics(telemetry)
-        pred_bits = logits.argmax(-1)
-        if self.decompress_output:
-            try:
-                pred_bits = model_output_decompress(pred_bits)
-            except Exception as e:
-                return {"error": f"Decompression failed: {e}", "suggestion": "Disable compression toggle."}
-        def _to_python(obj):
-            if isinstance(obj, torch.Tensor):
-                return obj.tolist()
-            if isinstance(obj, list):
-                return [_to_python(o) for o in obj]
-            if isinstance(obj, dict):
-                return {kk: _to_python(vv) for kk, vv in obj.items()}
-            return obj
-        tele = {k: _to_python(v) for k, v in telemetry.items()}
-        return {"predicted": pred_bits.squeeze(0).tolist(), "telemetry": tele, "ratio": ratio}
-
-    def infer_long(self, bits: torch.Tensor, ctx_bits: int = 4096, overlap: int = 256) -> Dict:
-        """Run sliding-window inference on a long sequence."""
-        assert self.model is not None
-        device = next(self.model.parameters()).device
-        bits = bits.to(device)
-        preds, logs = infer_long_sequence(self.model, bits.squeeze(0), ctx_bits=ctx_bits, overlap=overlap)
-        for tele in logs:
-            self._log_metrics(tele)
-        return {"predicted": preds.tolist(), "windows": len(logs)}
-
-    def _log_metrics(self, telemetry: Dict) -> None:
-        for key in self.metrics:
-            val = telemetry[key].mean().item()
-            self.metrics[key].append(val)
-        drift = detect_metric_drift(
-            self.metrics, window=self.drift_window, threshold=self.drift_threshold
-        )
-        bad = [k for k, v in drift.items() if v]
-        if bad:
-            warnings.warn(
-                f"Metric drift detected: {', '.join(bad)}",
-                MetricDriftWarning,
-            )
-
-    def infer_text(self, text: str) -> Dict[str, Any]:
-        """Run text through the model using the safety gate."""
-        assert self.model is not None
-        device = next(self.model.parameters()).device
-        bits = torch.tensor(text_to_bits(text), dtype=torch.long).unsqueeze(0).to(device)
-        out_bits, telemetry = hil_safe_inference(
-            self.model, bits, c_floor=self.c_floor, s_floor=self.s_floor
-        )
-        self._log_metrics(telemetry)
-        return {
-            "output": bits_to_text(out_bits.squeeze(0).tolist()),
-            "telemetry": telemetry,
-        }
-
-    def get_status(self) -> Dict[str, Any]:
-        info: Dict[str, Any] = {
-            "use_gpu": self.use_gpu,
-            "diffusion": self.diffusion,
-            "compression": self.use_compression,
-            "lambda_K": self.lambda_K,
-            "lambda_C": self.lambda_C,
-            "lambda_S": self.lambda_S,
-            "c_floor": self.c_floor,
-            "s_floor": self.s_floor,
-            "qat": self.qat,
-        }
-        if self.model is not None:
-            info.update(
-                {
-                    "d_model": self.model.d_model,
-                    "num_layers": self.model.num_layers,
-                    "d_ff": self.model.layers[0].linear1.out_features,
-                    "nhead": self.model.layers[0].self_attn.num_heads,
-                    "max_seq_len": self.model.pos_enc.pe.size(0),
-                }
-            )
-        else:
-            info.update(
-                {
-                    "d_model": None,
-                    "num_layers": 0,
-                    "d_ff": None,
-                    "nhead": None,
-                    "max_seq_len": None,
-                }
-            )
-        return info
-
-    def get_model_config(self) -> Dict[str, Any]:
-        """Return current model hyperparameters and safety settings."""
-        cfg: Dict[str, Any] = {
-            "lambda_K": self.lambda_K,
-            "lambda_C": self.lambda_C,
-            "lambda_S": self.lambda_S,
-            "c_floor": self.c_floor,
-            "s_floor": self.s_floor,
-        }
-        if self.model is not None:
-            cfg.update(
-                {
-                    "d_model": self.model.d_model,
-                    "nhead": self.model.layers[0].self_attn.num_heads,
-                    "num_layers": self.model.num_layers,
-                    "dim_feedforward": self.model.layers[0].linear1.out_features,
-                    "max_seq_len": self.model.pos_enc.pe.size(0),
-                    "chunk_size": self.model.chunk_size,
-                    "reversible": self.model.reversible,
-                    "use_checkpoint": self.model.use_checkpoint,
-                }
-            )
-        else:
-            cfg.update(
-                {
-                    "d_model": None,
-                    "nhead": None,
-                    "num_layers": 0,
-                    "dim_feedforward": None,
-                    "max_seq_len": None,
-                    "chunk_size": None,
-                    "reversible": None,
-                    "use_checkpoint": None,
-                }
-            )
-        return cfg
-
-    def get_metrics(self) -> Dict[str, Any]:
-        """Return logged telemetry metrics with summary statistics."""
-        from statistics import mean, stdev
-
-        data = {
-            "negentropy": self.metrics["negentropy_logits"],
-            "lz_complexity": self.metrics["lz_complexity_logits"],
-            "symbiosis": self.metrics["symbiosis_score"],
-        }
-        summary: Dict[str, Dict[str, float | None]] = {}
-        for key, values in data.items():
-            if values:
-                m = mean(values)
-                s = stdev(values) if len(values) > 1 else 0.0
-                summary[key] = {"mean": m, "std": s}
-            else:
-                summary[key] = {"mean": None, "std": None}
-        data["summary"] = summary
-        return data
-
-
-    def _save_state(self) -> None:
-        if self.model is None:
-            return
-        torch.save(self.model, self.weights_path)
-        with open(self.telemetry_log, "w") as f:
-            json.dump(self.metrics, f)
-
-
-manager: ModelManager | None = None
-
-
-@app.route("/")
-def index():
-    return render_template(
-        "dashboard.html",
-        metrics=manager.metrics,
-        lambdas={
-            "lambda_K": manager.lambda_K,
-            "lambda_C": manager.lambda_C,
-            "lambda_S": manager.lambda_S,
-        },
-        diffusion=manager.diffusion,
-        compression=manager.use_compression,
-        defaults={k: v.default for k, v in inspect.signature(BitTransformerLM.__init__).parameters.items() if v.default is not inspect._empty},
-        c_floor=manager.c_floor,
-        s_floor=manager.s_floor,
-        qat=manager.qat,
-    )
-
-
-@app.route("/status", methods=["GET"])
-def status():
-    if MCP_SERVER_ADDR:
-        return jsonify(mcp_get("/status"))
-    return jsonify(manager.get_status())
-
-
-@app.route("/model_config", methods=["GET"])
-def model_config():
-    if MCP_SERVER_ADDR:
-        return jsonify(mcp_get("/model_config"))
-    return jsonify(manager.get_model_config())
-
-
-@app.route("/metrics", methods=["GET"])
-def metrics():
-    if MCP_SERVER_ADDR:
-        return jsonify(mcp_get("/metrics"))
-    return jsonify(manager.get_metrics())
-
-
-@app.route("/save_checkpoint", methods=["POST"])
-def save_checkpoint_route():
-    repo_id = request.json.get("repo_id")
-    token = request.json.get("token") or os.getenv("HF_TOKEN")
-    if MCP_SERVER_ADDR:
-        return jsonify(mcp_post("/save_checkpoint", {"repo_id": repo_id, "token": token}))
-    if manager.model is None:
-        return jsonify({"error": "model not initialized"}), 400
-    if token:
-        hf_login(token=token)
-    save_checkpoint(manager.model, repo_id=repo_id)
-    return jsonify({"status": "saved"})
-
-
-@app.route("/download_checkpoint", methods=["POST"])
-def download_checkpoint_route():
-    repo_id = request.json.get("repo_id")
-    token = request.json.get("token") or os.getenv("HF_TOKEN")
-    if MCP_SERVER_ADDR:
-        return jsonify(mcp_post("/download_checkpoint", {"repo_id": repo_id, "token": token}))
-    if token:
-        hf_login(token=token)
-    dest = manager.weights_path + ".gz"
-    ok = download_checkpoint(dest, repo_id=repo_id)
-    if not ok:
-        return jsonify({"status": "failed"}), 500
-    if manager.model is None:
-        return jsonify({"status": "downloaded", "loaded": False})
-    with gzip.open(dest, "rb") as f:
-        state = torch.load(f, map_location="cpu")
-    manager.model.load_state_dict(state)
-    manager.optimizer, manager.scheduler = configure_optimizer(
-        manager.model, lr=1e-3, total_steps=manager.total_steps
-    )
-    manager._apply_device()
-    manager._save_state()
-    return jsonify({"status": "downloaded", "loaded": True})
-
-
-@app.route("/text_to_bits", methods=["POST"])
-def text_to_bits_route():
-    text = request.json.get("text", "")
-    if len(text) > 100_000:
-        return jsonify({"error": "text too large"}), 413
-    return jsonify({"bits": text_to_bits(text)})
-
-
-@app.route("/dataset", methods=["GET"])
-def dataset_route():
-    name = request.args.get("name", "")
-    split = request.args.get("split", "train")
-    size = int(request.args.get("size", 1))
-    seq_len = int(request.args.get("seq_len", 64))
-    if size * seq_len > 1_000_000:
-        return jsonify({"error": "dataset too large"}), 413
-    if name == "wikitext2":
-        try:
-            from datasets import load_dataset
-
-            ds = load_dataset("wikitext", "wikitext-2-raw-v1", split=split)
-            lines = [t for t in ds["text"] if t.strip()][:size]
-        except Exception:
-            bits = torch.randint(0, 2, (size, seq_len), dtype=torch.long)
-            return jsonify({"bits": bits.tolist()})
-        bits_list = []
-        for text in lines:
-            b = text_to_bits(text)[:seq_len]
-            if len(b) < seq_len:
-                b.extend([0] * (seq_len - len(b)))
-            bits_list.append(b)
-        if len(bits_list) < size:
-            pad = size - len(bits_list)
-            bits_list.extend(torch.randint(0, 2, (pad, seq_len), dtype=torch.long).tolist())
-        return jsonify({"bits": bits_list})
-    return jsonify({"error": "unknown dataset"}), 400
-
-
-@app.route("/init", methods=["POST"])
-def init_model():
-    data = request.json or {}
-    int_fields = {
-        "d_model",
-        "nhead",
-        "num_layers",
-        "dim_feedforward",
-        "max_seq_len",
-        "chunk_size",
-        "overlap",
-    }
-    float_fields = {"act_threshold"}
-    bool_fields = {"reversible", "use_checkpoint"}
-    params = {}
-    for k, v in data.items():
-        if v is None:
-            params[k] = None
-        elif k in int_fields:
-            params[k] = int(v)
-        elif k in float_fields:
-            params[k] = float(v)
-        elif k in bool_fields:
-            params[k] = bool(v)
-        else:
-            params[k] = v
-    if MCP_SERVER_ADDR:
-        data = mcp_post("/init", params)
-        return jsonify(data)
-    manager.init_model(params)
-    return jsonify({"status": "initialized", "params": params})
-
-
-@app.route("/train", methods=["POST"])
-def train_model():
-    bits = torch.tensor(request.json["bits"], dtype=torch.long)
-    if MCP_SERVER_ADDR:
-        data = mcp_post("/train", {"bits": request.json["bits"]})
-        return jsonify(data)
-    loss, ratio = manager.train_step(bits)
-    return jsonify({"loss": loss, "ratio": ratio})
-
-
-@app.route("/train_epochs", methods=["POST"])
-def train_epochs_route():
-    bits = torch.tensor(request.json["bits"], dtype=torch.long)
-    epochs = int(request.json.get("epochs", 1))
-    compress_prob = float(request.json.get("compress_prob", 0.5))
-    direct_prob = float(request.json.get("direct_prob", 0.0))
-    if MCP_SERVER_ADDR:
-        data = mcp_post(
-            "/train_epochs",
-            {
-                "bits": request.json["bits"],
-                "epochs": epochs,
-                "compress_prob": compress_prob,
-                "direct_prob": direct_prob,
-            },
-        )
-        return jsonify(data)
-    metrics = manager.train_epochs(
-        bits,
-        epochs=epochs,
-        compress_prob=compress_prob,
-        direct_prob=direct_prob,
-    )
-    return jsonify({"metrics": metrics})
-
-
-@app.route("/scale_up", methods=["POST"])
-def scale_up():
-    width_mult = float(request.json.get("width_mult", 1.0))
-    if MCP_SERVER_ADDR:
-        data = mcp_post("/scale_up", {"width_mult": width_mult})
-        return jsonify(data)
-    manager.scale_up(width_mult)
-    return jsonify({
-        "status": "scaled",
-        "layers": manager.model.num_layers,
-        "d_model": manager.model.d_model,
-    })
-
-
-@app.route("/collapse", methods=["POST"])
-def collapse_model():
-    cluster_bits = request.json["clusters"]
-    params = {k: int(v) for k, v in request.json["params"].items()}
-    width_scale = float(request.json.get("width_scale", 1.0))
-    if MCP_SERVER_ADDR:
-        data = mcp_post(
-            "/collapse",
-            {"clusters": cluster_bits, "params": params, "width_scale": width_scale},
-        )
-        return jsonify(data)
-    manager.collapse(cluster_bits, params, width_scale)
-    return jsonify({"status": "collapsed"})
-
-
-@app.route("/lambdas", methods=["GET", "POST"])
-def update_lambdas():
-    if request.method == "POST":
-        data = request.json
-        if MCP_SERVER_ADDR:
-            res = mcp_post("/lambdas", data)
-            return jsonify(res)
-        manager.set_lambdas(
-            float(data["lambda_K"]), float(data["lambda_C"]), float(data["lambda_S"])
-        )
-        return jsonify({"status": "updated"})
-    else:
-        if MCP_SERVER_ADDR:
-            return jsonify(mcp_get("/lambdas"))
-        return jsonify(
-            {
-                "lambda_K": manager.lambda_K,
-                "lambda_C": manager.lambda_C,
-                "lambda_S": manager.lambda_S,
-            }
-        )
-
-
-@app.route("/config/telemetry", methods=["GET", "POST"])
-def telemetry_config():
-    """Get or update telemetry λ weights and safety floors."""
-    if request.method == "POST":
-        data = request.json
-        if MCP_SERVER_ADDR:
-            res = mcp_post("/config/telemetry", data)
-            return jsonify(res)
-        manager.set_lambdas(
-            float(data.get("lambda_K", manager.lambda_K)),
-            float(data.get("lambda_C", manager.lambda_C)),
-            float(data.get("lambda_S", manager.lambda_S)),
-        )
-        manager.set_floors(
-            float(data.get("c_floor", manager.c_floor)),
-            float(data.get("s_floor", manager.s_floor)),
-        )
-        return jsonify({"status": "updated"})
-    else:
-        if MCP_SERVER_ADDR:
-            return jsonify(mcp_get("/config/telemetry"))
-        return jsonify(
-            {
-                "lambda_K": manager.lambda_K,
-                "lambda_C": manager.lambda_C,
-                "lambda_S": manager.lambda_S,
-                "c_floor": manager.c_floor,
-                "s_floor": manager.s_floor,
-            }
-        )
-
-
-@app.route("/diffusion", methods=["GET", "POST"])
-def update_diffusion():
-    if request.method == "POST":
-        if MCP_SERVER_ADDR:
-            return jsonify(mcp_post("/diffusion", request.json))
-        manager.set_diffusion(bool(request.json.get("diffusion", False)))
-        return jsonify({"status": "updated"})
-    else:
-        if MCP_SERVER_ADDR:
-            return jsonify(mcp_get("/diffusion"))
-        return jsonify({"diffusion": manager.diffusion})
-
-
-@app.route("/gpu", methods=["GET", "POST"])
-def update_gpu():
-    if request.method == "POST":
-        if MCP_SERVER_ADDR:
-            return jsonify(mcp_post("/gpu", request.json))
-        manager.set_gpu(bool(request.json.get("use_gpu", False)))
-        return jsonify({"status": "updated"})
-    else:
-        if MCP_SERVER_ADDR:
-            return jsonify(mcp_get("/gpu"))
-        return jsonify({"use_gpu": manager.use_gpu})
-
-
-@app.route("/compression", methods=["GET", "POST"])
-def update_compression():
-    if request.method == "POST":
-        if MCP_SERVER_ADDR:
-            return jsonify(mcp_post("/compression", request.json))
-        manager.set_compression(bool(request.json.get("compression", False)))
-        return jsonify({"status": "updated"})
-    else:
-        if MCP_SERVER_ADDR:
-            return jsonify(mcp_get("/compression"))
-        return jsonify({"compression": manager.use_compression})
-
-
-@app.route("/qat", methods=["GET", "POST"])
-def update_qat():
-    if request.method == "POST":
-        if MCP_SERVER_ADDR:
-            return jsonify(mcp_post("/qat", request.json))
-        manager.set_qat(bool(request.json.get("qat", False)))
-        return jsonify({"status": "updated"})
-    else:
-        if MCP_SERVER_ADDR:
-            return jsonify(mcp_get("/qat"))
-        return jsonify({"qat": manager.qat})
-
-
-@app.route("/infer", methods=["POST"])
-def inference():
-    bits = torch.tensor(request.json["bits"], dtype=torch.long)
-    if MCP_SERVER_ADDR:
-        data = mcp_post("/infer", {"bits": request.json["bits"]})
-        return jsonify(data)
-    result = manager.infer(bits)
-    return jsonify(result)
-
-
-@app.route("/infer_long", methods=["POST"])
-def inference_long():
-    bits = torch.tensor(request.json["bits"], dtype=torch.long)
-    ctx = int(request.json.get("ctx_bits", 4096))
-    overlap = int(request.json.get("overlap", 256))
-    if MCP_SERVER_ADDR:
-        data = mcp_post(
-            "/infer_long",
-            {"bits": request.json["bits"], "ctx_bits": ctx, "overlap": overlap},
-        )
-        return jsonify(data)
-    result = manager.infer_long(bits, ctx_bits=ctx, overlap=overlap)
-    return jsonify(result)
-
-
-@app.route("/infer_text", methods=["POST"])
-def inference_text():
-    text = request.json.get("text", "")
-    if MCP_SERVER_ADDR:
-        data = mcp_post("/infer_text", {"text": text})
-        return jsonify(data)
-    result = manager.infer_text(text)
-    return jsonify(result)
-
-@app.route("/plot.png")
-def plot_png():
-    if MCP_SERVER_ADDR:
-        resp = requests.get(MCP_SERVER_ADDR.rstrip("/") + "/plot.png")
-        resp.raise_for_status()
-        return send_file(io.BytesIO(resp.content), mimetype="image/png")
-    fig, _ = plot_telemetry(manager.metrics)
-    buf = io.BytesIO()
-    fig.savefig(buf, format="png")
-    plt.close(fig)
-    buf.seek(0)
-    return send_file(buf, mimetype="image/png")
-
-
-def run_dashboard(host: str | None = None, port: int | None = None,
-                  snapshot_dir: str | None = None, telemetry_log: str | None = None) -> None:
-    """Launch the Flask dashboard server."""
-    env_host = os.getenv("HOST", "0.0.0.0")
-    env_port = int(os.getenv("PORT", "5000"))
-    host = host or env_host
-    port = port or env_port
-    global manager
-    if manager is None:
-        manager = ModelManager(snapshot_dir, telemetry_log)
-    app.run(host=host, port=port, debug=True)
-
-
-if __name__ == "__main__":
-    import argparse
-
-    parser = argparse.ArgumentParser(description="Run dashboard server")
-    parser.add_argument("--host", default=os.getenv("HOST", "0.0.0.0"))
-    parser.add_argument("--port", type=int, default=int(os.getenv("PORT", "5000")))
-    parser.add_argument("--snapshot-dir", default=os.getenv("SNAPSHOT_DIR", "snapshots"))
-    parser.add_argument("--telemetry-log", default=os.getenv("TELEMETRY_LOG"))
-    args = parser.parse_args()
-    run_dashboard(args.host, args.port, args.snapshot_dir, args.telemetry_log)