app.py

"""
DaisyChain — interactive routing demo (HuggingFace Space).

Paste DNA; the learned router reads how *surprised* each ~74M specialist is (bits/base)
plus its hidden state and hands the sequence to its home specialist — then that specialist
streams a continuation live. Styled after the Modular-Mind panel: animated routing cards,
a first-run loading notice, live token streaming. Every handler is a generator.
"""
import html as _h
import os
import json

import gradio as gr

# ZeroGPU: @spaces.GPU allocates a GPU only for the decorated call. Falls back to a no-op
# decorator when `spaces` isn't installed (local / plain CPU).
try:
    import spaces
    _gpu = spaces.GPU
except Exception:
    def _gpu(fn=None, **kw):
        return fn if callable(fn) else (lambda f: f)

from daisychain import DaisyChain

HERE = os.path.dirname(os.path.abspath(__file__))
MODEL_REPO = os.environ.get("DAISYCHAIN_REPO", "DaisyChainAI/daisychain-genomics")
DEVICE = os.environ.get("DAISYCHAIN_DEVICE", "cpu")

# code + tokenizer + router are bundled here; pull the big specialist weights from the
# (gated) model repo on first launch using the HF_TOKEN Space secret. No silent
# swallow — if the download fails we want a visible error, not a broken-but-running app.
if not os.path.exists(os.path.join(HERE, "eukaryote", "model.safetensors")):
    from huggingface_hub import snapshot_download
    snapshot_download(MODEL_REPO, local_dir=HERE,
                      token=os.environ.get("HF_TOKEN"),
                      allow_patterns=["*/model.safetensors", "tokenizer.json", "router2.pt"])

_DC = {"m": None}            # lazy-loaded so CUDA is never touched at import
_WARMED = {"done": False}    # so the "loading" notice only shows on the first run

EMOJI = {"eukaryote": "🧬 Eukaryote", "prokaryote": "🦠 Prokaryote",
         "mrna": "📜 mRNA", "mrna_splice": "✂️ mRNA-splice"}
COLOR = {"eukaryote": "#7c5cff", "prokaryote": "#22b8cf",
         "mrna": "#e64980", "mrna_splice": "#37b24d"}
DESC = DaisyChain.DESCRIPTIONS


def _moe():
    if _DC["m"] is None:
        _DC["m"] = DaisyChain(root=HERE, device=DEVICE)
    return _DC["m"]


# ---- HTML rendering (ported from the Modular-Mind panel) --------------------------
_CSS = """<style>
.dcx{font-family:system-ui,sans-serif;color:#dde;margin:4px 0}
.dcx .note{background:#14141c;border:1px solid #2a2a35;border-radius:10px;padding:12px 14px;color:#9bd;font-size:14px}
.dcx .h{font-size:17px;font-weight:800;margin:4px 0 8px}
.dcx .p{color:#8892a8}
.dcx .g{color:#eef2ff;font-weight:600}
.dcx .chain{display:flex;gap:8px;align-items:stretch;flex-wrap:wrap;margin:6px 0}
.dcx .link{align-self:center;color:#445;font-size:20px;margin-bottom:18px}
.dcx .card{flex:1;min-width:190px;background:#14141c;border:1px solid #2a2a35;border-radius:12px;padding:11px 13px;position:relative;overflow:hidden}
.dcx .card .nm{font-weight:800;font-size:15px}
.dcx .card .meta{color:#99a;font-size:11px;margin-top:2px;min-height:26px}
.dcx .card .bar{height:10px;background:#23232e;border-radius:6px;margin-top:8px;overflow:hidden}
.dcx .card .fill{height:100%;border-radius:6px;animation:dcxw .7s ease}
.dcx .card .pct{font-size:12px;color:#bcd;margin-top:4px}
.dcx .badge{position:absolute;top:9px;right:10px;font-size:10px;font-weight:800;letter-spacing:.08em;padding:3px 8px;border-radius:99px;color:#0a1410}
@keyframes dcxw{from{width:0}}
.dcx .gen{background:#101018;border:1px solid #2a2a35;border-radius:12px;padding:13px 15px;margin:10px 0;font-size:15px;line-height:1.7;font-family:ui-monospace,SFMono-Regular,Menlo,Consolas,monospace;word-break:break-all}
.dcx .caret{display:inline-block;width:9px;height:17px;border-radius:2px;background:#7ad1ff;margin-left:2px;vertical-align:text-bottom;animation:dcxb .8s steps(1) infinite}
@keyframes dcxb{50%{opacity:0}}
.dcx .sub{color:#889;font-size:12px;line-height:1.5;margin-top:8px}
.dcx .stats{display:flex;gap:10px;flex-wrap:wrap;margin:10px 0}
.dcx .stat{flex:1;min-width:130px;text-align:center;background:#14141c;border:1px solid #2a2a35;border-radius:12px;padding:13px 8px}
.dcx .stat .v{font-size:28px;font-weight:800;line-height:1}
.dcx .stat .l{font-size:11px;color:#99a;margin-top:6px}
.dcx table{border-collapse:collapse;width:100%;margin:8px 0;font-size:13.5px}
.dcx th,.dcx td{border:1px solid #2a2a35;padding:7px 10px;text-align:left}
.dcx th{background:#14141c;color:#bcd}
.dcx td.n{text-align:right;font-variant-numeric:tabular-nums}
</style>"""


def _wrap(body):
    return _CSS + "<div class='dcx'>" + body + "</div>"


def _esc(s):
    return _h.escape(s or "").replace("\n", "<br>")


def _notice(action="Routing"):
    if not _WARMED["done"]:
        try:
            gr.Info("First run — loading the four ~74M specialists (~20–40s on CPU). After this it's quick.")
        except Exception:
            pass
        return _wrap(f"<div class='note'>⏳ Loading the four ~74M specialists + {action.lower()}… "
                     "first run can take ~20–40s on CPU; every run after is fast.</div>")
    return _wrap(f"<div class='note'>⏳ {action}…</div>")


def _msg(title, body):
    return _wrap(f"<div class='note'><b>{title}</b><br>{body}</div>")


def _cards(bpb, winner=None):
    """One animated card per specialist: surprise (bits/base), confidence bar, winner badge + glow.
    bpb values may be None (not computed yet). Lower bits/base = more 'at home' = fuller bar."""
    cells = []
    doms = list(bpb.keys())
    for i, n in enumerate(doms):
        c = COLOR.get(n, "#9b59b6")
        v = bpb[n]
        win = (n == winner)
        conf = max(0.0, min(1.0, (2.02 - v) / 0.5)) if v is not None else 0.0  # ~1.52..2.02 -> 1..0
        style = f"border-color:{c};box-shadow:0 0 16px {c}40" if win else ""
        badge = f"<span class='badge' style='background:{c}'>ROUTED ✓</span>" if win else ""
        meta = (f"{DESC.get(n,'')}<br>{v:.3f} bits/base (lower = more at home)"
                if v is not None else f"{DESC.get(n,'')}<br>…")
        bar = (f"<div class='bar'><div class='fill' style='width:{conf*100:.1f}%;background:{c}'></div></div>"
               f"<div class='pct'>confidence {conf*100:.0f}%</div>") if v is not None else \
              "<div class='bar'></div><div class='pct'>…</div>"
        cells.append(
            f"<div class='card' style='{style}'>{badge}"
            f"<div class='nm' style='color:{c}'>{EMOJI.get(n, n)}</div>"
            f"<div class='meta'>{meta}</div>{bar}</div>")
        if i < len(doms) - 1:
            cells.append("<div class='link'>⬭</div>")
    return "<div class='chain'>" + "".join(cells) + "</div>"


def _gen_box(prompt, gen, live=False):
    caret = "<span class='caret'></span>" if live else ""
    return (f"<div class='gen'><span class='p'>{_esc(prompt)}</span>"
            f"<span class='g'>{_esc(gen)}</span>{caret}</div>")


FOOTER = ("Four ~74M DNA/RNA specialists (≈295M total, <b>under Carbon-500M</b>), each distilled "
          "per-domain from Carbon-500M. A learned router reads every specialist's surprise + hidden "
          "state and routes to the home specialist — held-out routing accuracy <b>94.8%</b>. Only one "
          "specialist runs per query (~7× cheaper than the 500M monolith).")


# ---- handler ----------------------------------------------------------------------
@_gpu(duration=120)
def route_run(seq, n_bases, do_gen):
    yield _notice("Routing & generating")
    seq = (seq or "").strip()
    if len(seq) < 18:
        yield _msg("🧬 Enter a DNA sequence", "Paste at least 18 bases (A/C/G/T) — try an example below.")
        return
    dc = _moe()
    doms = dc.domains
    bpb = {d: None for d in doms}
    # progressively reveal each specialist's surprise (the chain lighting up)
    sc, hd = dc._scores_hidden(seq)
    for d in doms:
        bpb[d] = sc[d] / 6 / 0.6931
        yield _wrap("<div class='h'>🔗 Sending the sequence down the chain…</div>" + _cards(bpb))
    home, _ = dc.route(seq)
    c = COLOR.get(home, "#9b59b6")
    head = (f"<div class='h'>🧭 Routed to <span style='color:{c}'>{EMOJI.get(home, home)}</span>"
            f" — the specialist most at home with your sequence</div>" + _cards(bpb, winner=home))
    if do_gen:
        for gen in dc.generate_stream(home, length=int(n_bases), temperature=0.9, top_k=20, prompt=seq[-60:]):
            yield _wrap(head + _gen_box(seq[-60:], gen, live=True))
        _WARMED["done"] = True
        yield _wrap(head + _gen_box(seq[-60:], gen, live=False) + f"<div class='sub'>{FOOTER}</div>")
    else:
        _WARMED["done"] = True
        yield _wrap(head + f"<div class='sub'>{FOOTER}</div>")


STATS_HTML = _wrap(
    "<div class='h'>📊 DaisyChain vs Carbon-500M — the fair baseline</div>"
    "<div class='stats'>"
    "<div class='stat'><div class='v' style='color:#37b24d'>94.8%</div>"
    "<div class='l'>routing accuracy<br>(held-out)</div></div>"
    "<div class='stat'><div class='v' style='color:#7c5cff'>≈295M</div>"
    "<div class='l'>total params<br>(4 × ~74M) &lt; Carbon-500M</div></div>"
    "<div class='stat'><div class='v' style='color:#22b8cf'>~7×</div>"
    "<div class='l'>cheaper per query<br>(one 74M specialist active)</div></div>"
    "</div>"
    "<table><tr><th>metric</th><th>DaisyChain</th><th>Carbon-500M</th></tr>"
    "<tr><td>Likelihood — bits/base (↓ better)</td><td class='n'>1.86</td><td class='n'>1.75</td></tr>"
    "<tr><td>Seq-recovery, eukaryote (↑ better)</td><td class='n'>31.8%</td><td class='n'>42.2%</td></tr>"
    "<tr><td>Seq-recovery, bacteria (↑ better)</td><td class='n'>34.0%</td><td class='n'>49.5%</td></tr>"
    "</table>"
    "<div class='sub'>Four ~74M specialists (≈295M total, <b>under Carbon-500M</b>); only one runs per "
    "query, so it's ~7× cheaper per token. Behind the 500M / 1T-token monolith but within striking "
    "distance — the gap is concentrated in the structured domains (mRNA, bacteria) and keeps closing "
    "with more per-domain training. Same protocols as Carbon's eval suite (sequence recovery; per-base "
    "likelihood). Carbon-500M is the right yardstick for a sub-500M modular set, not the 3B flagship.</div>")


HERO = """# 🌼 DaisyChain — a modular genomic mind
**Four ~74M DNA/RNA specialists (≈295M total, under Carbon-500M)** — 🧬 Eukaryote, 🦠 Prokaryote,
📜 mRNA, ✂️ mRNA-splice — each **distilled per-domain from Carbon-500M**. A learned router reads how
*surprised* each specialist is by your sequence (bits/base) plus its hidden state, and hands the work
to its **home specialist**. Paste DNA and watch it route in real time.

> ℹ️ *Research demo: tiny specialists trained on a slice of the Carbon corpus — the **routing** (which
> specialist is most at home) and the **sub-500M modular architecture** are the point, not Carbon-level
> generation.*"""


def build():
    with gr.Blocks(title="DaisyChain — modular genomic mind", theme=gr.themes.Soft()) as demo:
        with gr.Accordion("🌼 DaisyChain — independent DNA specialists behind a learned router", open=True):
            gr.Markdown(HERO)
            with gr.Row():
                seq = gr.Textbox(label="DNA sequence", lines=3, scale=4,
                                 placeholder="ACGT… (eukaryotic, bacterial, mRNA, or splice-site DNA)")
                n = gr.Slider(60, 300, value=150, step=30, label="generate bases", scale=1)
            with gr.Row():
                gen_ck = gr.Checkbox(value=True, label="stream a continuation from the routed specialist")
                btn = gr.Button("🔗 Route through the DaisyChain", variant="primary")
            out = gr.HTML(_wrap(_cards({d: None for d in DaisyChain.DESCRIPTIONS})))
            btn.click(route_run, [seq, n, gen_ck], out)
            try:
                ex = json.load(open(os.path.join(HERE, "examples.json")))
                gr.Examples([[v, 150, True] for v in ex.values()], inputs=[seq, n, gen_ck],
                            label="Example sequences (one per domain)")
            except Exception:
                pass
            gr.HTML(STATS_HTML)
    return demo


if __name__ == "__main__":
    build().launch()