File size: 3,103 Bytes
9409aaa 0f66bac 9409aaa cd53ece 9409aaa 0f66bac 9409aaa cd53ece 9409aaa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | """
space/server_app.py — One for All, custom frontend on gr.Server.
gr.Server = FastAPI with Gradio's API engine: @app.api functions become
queued/streaming endpoints (SSE), and we serve our own HTML/JS at "/" —
the browser talks to the endpoints via @gradio/client (required for ZeroGPU).
Run locally:
cd space && VIZ_DATA_PATH=/path/to/viz_data.json python server_app.py
The legacy Blocks UI (app.py) keeps working — both share _boot.py.
"""
from __future__ import annotations
from pathlib import Path
import spaces
from fastapi.responses import FileResponse, HTMLResponse
from gradio import Server
import _boot
RT = _boot.load_runtime()
FRONTEND = Path(__file__).parent / "frontend"
app = Server()
# ── API: static viz payload (soul space, CKA, curves, meta) ────────────────
@app.api(name="viz")
def viz() -> dict:
return _boot.viz_payload(RT.viz, RT.coords3d, RT.backend, RT.model_ready)
# ── API: streaming probe — text + live gates, final soul-space point ───────
@app.api(name="probe", stream_every=0.1)
@spaces.GPU
def probe(text: str) -> dict: # annotation = type of each streamed chunk
if not text.strip():
yield {"error": "empty prompt", "done": True}
return
_boot.to_device(RT) # lazy-loads the student here (real CUDA inside @spaces.GPU)
if not RT.model_ready:
yield {"error": RT.load_error or "model not ready", "done": True}
return
partial, gates = "", []
for partial, gates in _boot.stream(RT, text):
yield {"text": partial, "gates": gates, "done": False}
point = None
if RT.reducer is not None:
point, gates = _boot.final_probe(RT, text)
yield {"text": partial, "gates": gates, "point": point, "done": True}
# ── API: streaming arena — base (LoRA off) vs deku, interleaved ────────────
@app.api(name="arena", stream_every=0.1)
@spaces.GPU
def arena(text: str) -> dict: # annotation = type of each streamed chunk
if not text.strip():
yield {"error": "empty prompt", "done": True}
return
_boot.to_device(RT) # lazy-loads the student here (real CUDA inside @spaces.GPU)
if not RT.model_ready:
yield {"error": RT.load_error or "model not ready", "done": True}
return
base = deku = ""
gates: list[float] = []
for base, deku, gates in _boot.stream_pair(RT, text):
yield {"base": base, "deku": deku, "gates": gates, "done": False}
yield {"base": base, "deku": deku, "gates": gates, "done": True}
# ── Custom frontend (overrides Gradio's default UI at "/") ────────────────
@app.get("/", response_class=HTMLResponse)
async def home():
return (FRONTEND / "index.html").read_text()
@app.get("/static/{name}")
async def static(name: str):
target = (FRONTEND / name).resolve()
if target.parent != FRONTEND.resolve() or not target.exists():
return HTMLResponse("not found", status_code=404)
return FileResponse(target)
if __name__ == "__main__":
app.launch()
|