Spaces:
Running
Running
| """Local Model Explorer: FastAPI backend + static frontend.""" | |
| from __future__ import annotations | |
| import dataclasses | |
| import json | |
| import logging | |
| import os | |
| import threading | |
| import time | |
| from contextlib import asynccontextmanager | |
| from pathlib import Path | |
| from typing import Literal | |
| from fastapi import FastAPI, HTTPException, Request | |
| from fastapi.exceptions import RequestValidationError | |
| from fastapi.responses import FileResponse, JSONResponse | |
| from fastapi.staticfiles import StaticFiles | |
| from pydantic import BaseModel, Field, ValidationError, model_validator | |
| from .catalogue import MLX_EXPLORER, ONNX_EXPLORER, ROOT, Catalogue | |
| from .detail import BUDGET, REPOS_ON_OPEN, REPOS_PER_MODEL, DetailCache | |
| from .events import EventBatch, to_row | |
| from .hardware import CONTEXTS, KV_TYPES, UNIFIED, Hardware, presets | |
| from .parsing import SIZE_BUCKETS | |
| from .ratelimit import RateLimiter | |
| from .recommend import PRIORITIES, SORTS, WEIGHTS, Engine, Query, options_for, sort_results | |
| from .sink import make_sink | |
| from .stats import StatsCache | |
| logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s") | |
| log = logging.getLogger("explorer") | |
| logging.getLogger("uvicorn.access").disabled = True # never log client addresses | |
| logging.getLogger("httpx").setLevel(logging.WARNING) | |
| APP_VERSION = os.environ.get("EXPLORER_VERSION", "0.1.0") | |
| STATIC = ROOT / "static" | |
| MAX_BODY = 64 * 1024 | |
| FRAME_ANCESTORS = os.environ.get("EXPLORER_FRAME_ANCESTORS", "'self' https://huggingface.co https://*.hf.space") | |
| DETAIL_TOP_N = int(os.environ.get("EXPLORER_DETAIL_TOP_N", "12")) | |
| PREWARM_N = int(os.environ.get("EXPLORER_PREWARM", "400")) | |
| RELOAD_SECONDS = float(os.environ.get("EXPLORER_RELOAD_SECONDS", "900")) | |
| class HardwareIn(BaseModel): | |
| kind: Literal["gpu", "unified", "cpu"] | |
| gpus: list[float] = Field(default_factory=list, max_length=16) | |
| ram_gb: float = Field(32, gt=0, le=4096) | |
| unified_system: str | None = None | |
| def _check(self): | |
| if any(not (0 < v <= 256) for v in self.gpus): | |
| raise ValueError("each GPU's VRAM must be between 0 and 256 GB") | |
| if self.kind == "gpu" and not self.gpus: | |
| raise ValueError("a GPU setup needs at least one GPU") | |
| if self.unified_system is not None and self.unified_system not in UNIFIED: | |
| raise ValueError("unknown unified-memory system") | |
| return self | |
| def to_hardware(self) -> Hardware: | |
| share = UNIFIED[self.unified_system][1] if self.kind == "unified" and self.unified_system else None | |
| return Hardware(self.kind, tuple(self.gpus) if self.kind == "gpu" else (), self.ram_gb, share) | |
| class RecommendRequest(BaseModel): | |
| hardware: HardwareIn | None = None | |
| context: int = 8192 | |
| kv_type: Literal["f16", "q8_0", "q4_0"] = "f16" | |
| priority: Literal["balanced", "quality", "small"] = "balanced" | |
| min_bits: Literal[0, 2, 3, 4, 5, 6, 8] = 0 | |
| vram_only: bool = False | |
| mmproj: Literal["gpu", "ram", "off"] = "gpu" | |
| sort: Literal["best", "popular", "recent", "community"] = "best" | |
| family: str | None = Field(None, max_length=40) | |
| size: str | None = Field(None, max_length=10) | |
| search: str | None = Field(None, max_length=80) | |
| moe_only: bool = False | |
| vision_only: bool = False | |
| fits_only: bool = False | |
| offset: int = Field(0, ge=0, le=100_000) | |
| limit: int = Field(24, ge=1, le=60) | |
| class ModelRequest(BaseModel): | |
| key: str = Field(..., max_length=200) | |
| hardware: HardwareIn | None = None | |
| context: int = 8192 | |
| kv_type: Literal["f16", "q8_0", "q4_0"] = "f16" | |
| priority: Literal["balanced", "quality", "small"] = "balanced" | |
| min_bits: Literal[0, 2, 3, 4, 5, 6, 8] = 0 | |
| vram_only: bool = False | |
| mmproj: Literal["gpu", "ram", "off"] = "gpu" | |
| class State: | |
| def __init__(self, catalogue=None, details=None, sink=None, fetch_live=True): | |
| base = Path(os.environ.get("EXPLORER_DATA_DIR", ROOT / ".runtime")) | |
| self.catalogue = catalogue or Catalogue() | |
| self.details = details or DetailCache() | |
| self.sink = sink or make_sink(base) | |
| self.engine = Engine() | |
| self.limiter = RateLimiter(rate_per_minute=float(os.environ.get("EXPLORER_RATE_PER_MIN", "120")), burst=60) | |
| self.existing_rows: list[dict] = [] | |
| self.stats = StatsCache(lambda: self.existing_rows + self.sink.live_rows()) | |
| self.fetch_live = fetch_live | |
| self.collection_enabled = os.environ.get("EXPLORER_SINK", "local") != "off" | |
| def prewarm(self): | |
| """File listings (not file downloads) for the most popular models, gently.""" | |
| top = sorted(self.catalogue.groups.values(), key=lambda g: -g.downloads)[:PREWARM_N] | |
| for g in top: | |
| try: | |
| if any(self.details._stale(r.id) for r in g.repos[:REPOS_PER_MODEL]): | |
| self.details.get(g, timeout=120, with_geometry=False) | |
| time.sleep(2.0) # leave most of the Hub budget for people using the page | |
| except Exception: | |
| pass | |
| def boot(self): | |
| self.catalogue.load_offline() | |
| if self.fetch_live: | |
| def initial(): | |
| self.catalogue.refresh() | |
| self.prewarm() | |
| threading.Thread(target=initial, daemon=True, name="catalogue-initial").start() | |
| self.catalogue.start_background_refresh() | |
| self.sink.start() | |
| def reload_loop(): | |
| for delay in (0, 90): | |
| time.sleep(delay) | |
| self.reload_existing() | |
| while True: | |
| time.sleep(RELOAD_SECONDS) | |
| self.reload_existing() | |
| threading.Thread(target=reload_loop, daemon=True, name="reload-existing").start() | |
| def reload_existing(self): | |
| started = time.time() | |
| rows = self.sink.read_existing() | |
| if rows is None: | |
| return | |
| self.existing_rows = rows | |
| self.sink.forget_flushed_before(started) | |
| self.stats.invalidate() | |
| def _hidden(s, details) -> bool: | |
| return _only_auxiliary(s) or s.group.key in details.not_llama_cpp | |
| def _only_auxiliary(s) -> bool: | |
| """Every file we could read is too small to be the model (a repo of draft models or | |
| speculative-decoding heads, which nobody runs on their own), or below the chosen bit level.""" | |
| repos = (s.detail_repos or {}).values() | |
| return s.detail_loaded and not s.options and bool(repos) and all(rf.error is None for rf in repos) | |
| def _query(hw: HardwareIn | None, context: int, kv_type: str, priority: str, min_bits: int = 0, | |
| vram_only: bool = False, mmproj: str = "gpu") -> Query: | |
| if context not in CONTEXTS: | |
| raise HTTPException(422, "context must be one of " + ",".join(map(str, CONTEXTS))) | |
| return Query(hw.to_hardware() if hw else None, context, kv_type, priority, min_bits, | |
| vram_only and bool(hw) and hw.kind == "gpu", mmproj) | |
| def create_app(state: State | None = None) -> FastAPI: | |
| st = state or State() | |
| async def lifespan(app): | |
| if state is None: | |
| st.boot() | |
| yield | |
| st.sink.flush() | |
| app = FastAPI(title="Local Model Explorer", version=APP_VERSION, lifespan=lifespan, | |
| docs_url=None, redoc_url=None, openapi_url=None) | |
| app.state.st = st | |
| async def guard(request: Request, call_next): | |
| cl = request.headers.get("content-length") | |
| if cl and cl.isdigit() and int(cl) > MAX_BODY: | |
| return JSONResponse({"error": "payload too large"}, status_code=413) | |
| resp = await call_next(request) | |
| resp.headers["X-Content-Type-Options"] = "nosniff" | |
| resp.headers["Referrer-Policy"] = "no-referrer" | |
| resp.headers["Content-Security-Policy"] = ( | |
| "default-src 'self'; img-src 'self' data:; style-src 'self'; script-src 'self'; " | |
| f"connect-src 'self'; frame-ancestors {FRAME_ANCESTORS}") | |
| return resp | |
| async def _invalid(request, exc): | |
| errs = [{"loc": [str(x) for x in e.get("loc", [])][-3:], "msg": str(e.get("msg"))[:160]} | |
| for e in exc.errors()[:10]] | |
| return JSONResponse({"error": "invalid request", "details": errs}, status_code=422) | |
| def client_key(request: Request) -> str: | |
| fwd = request.headers.get("x-forwarded-for", "") | |
| return fwd.split(",")[0].strip() or (request.client.host if request.client else "unknown") | |
| def index(): | |
| return FileResponse(STATIC / "index.html") | |
| def stats_page(): | |
| return FileResponse(STATIC / "stats.html") | |
| def healthz(): | |
| cat, sink = st.catalogue.status(), st.sink.status() | |
| ok = cat["groups"] > 0 | |
| return JSONResponse({"ok": ok, "version": APP_VERSION, "catalogue": cat, "sink": sink, | |
| "geometries": len(st.details.geometries)}, status_code=200 if ok else 503) | |
| def meta(): | |
| return { | |
| "version": APP_VERSION, "engine": st.engine.version, "families": st.catalogue.families(), | |
| "sizes": SIZE_BUCKETS, "priorities": PRIORITIES, "sorts": SORTS, "score_weights": WEIGHTS, "hardware": presets(), | |
| "catalogue": st.catalogue.status(), "collection_enabled": st.collection_enabled, | |
| "links": {"mlx": MLX_EXPLORER, "onnx": ONNX_EXPLORER}, | |
| } | |
| def recommend(req: RecommendRequest): | |
| q = _query(req.hardware, req.context, req.kv_type, req.priority, req.min_bits, req.vram_only, req.mmproj) | |
| groups = st.catalogue.matching(req.family or None, req.size or None, req.search, req.moe_only, req.vision_only) | |
| community = st.stats.community() | |
| cached = {g.key: d for g in groups if (d := st.details.peek(g))} | |
| prelim = [s for s in st.engine.score(groups, q, cached, community) if not _hidden(s, st.details)] | |
| if (req.fits_only or q.vram_only) and q.hardware: | |
| prelim = [s for s in prelim if s.best and s.best.fit and s.best.fit.tier not in ("no", "no_vram")] | |
| prelim = sort_results(prelim, req.sort, community) | |
| # Read files for everything on the page (and a little beyond it on the first page, where | |
| # exact sizes can reorder the top results); reads that miss the wait keep going. | |
| window = prelim[: max(DETAIL_TOP_N, req.limit)] if req.offset == 0 else prelim[req.offset: req.offset + req.limit] | |
| # Exact file sizes and geometry for what is about to be shown. | |
| missing = [s.group for s in window if not s.detail_loaded or s.best is None or s.best.estimated] | |
| if missing and st.fetch_live: | |
| futs = [st.details.pool.submit(st.details.get, g, 4.0) for g in missing] | |
| deadline = time.time() + 4.0 | |
| for f in futs: | |
| try: | |
| f.result(timeout=max(deadline - time.time(), 0.05)) | |
| except Exception: | |
| pass | |
| fresh = {g.key: d for g in missing if (d := st.details.peek(g))} | |
| rescored = {s.group.key: s for s in st.engine.score([s.group for s in window], q, {**cached, **fresh}, community)} | |
| prelim = [rescored.get(s.group.key, s) for s in prelim] | |
| prelim = [s for s in prelim if not _hidden(s, st.details)] | |
| if (req.fits_only or q.vram_only) and q.hardware: | |
| prelim = [s for s in prelim if s.best and s.best.fit and s.best.fit.tier not in ("no", "no_vram")] | |
| prelim = sort_results(prelim, req.sort, community) | |
| page = prelim[req.offset: req.offset + req.limit] | |
| notices = [] | |
| if st.catalogue.degraded: | |
| notices.append("Showing a cached copy of the model list.") | |
| # Only models whose files are still to be read; a model whose GGUFs can't be planned stays marked (est.) | |
| est = sum(1 for s in page if s.best and s.best.estimated and s.group.key not in st.details.geometry_failed) | |
| if est: | |
| notices.append(f"{est} result(s) are estimated from the parameter count until their files are read.") | |
| return {"total": len(prelim), "offset": req.offset, "results": [s.to_dict(q) for s in page], | |
| "notices": notices, "engine": st.engine.version} | |
| def model(req: ModelRequest): | |
| g = st.catalogue.groups.get(req.key) | |
| if g is None: | |
| raise HTTPException(404, "model not found") | |
| q = _query(req.hardware, req.context, req.kv_type, req.priority, req.min_bits, req.vram_only, req.mmproj) | |
| detail = st.details.get(g, timeout=25.0, repos=REPOS_ON_OPEN) | |
| community = st.stats.community() | |
| scored = st.engine.score([g], q, {g.key: detail}, community)[0] | |
| # The recommendation first, then what runs best, then higher quality, then more downloaded uploaders. | |
| best = scored.best | |
| options = sorted(scored.options, key=lambda o: (o is not best, -(o.fit.score if o.fit else 0), | |
| -(o.quant.bpw or 0), -o.repo_downloads, o.repo)) | |
| repos = [{**r.to_dict(), "files_error": (detail.repos.get(r.id).error if detail.repos.get(r.id) else None), | |
| "loaded": r.id in detail.repos} for r in g.repos] | |
| other = {fmt: ids[:30] for fmt, ids in g.other_formats.items()} | |
| below = (len(options_for(g, detail, dataclasses.replace(q, min_bits=0))) - len(options)) if q.min_bits else 0 | |
| return { | |
| **scored.to_dict(q), | |
| "options": [o.to_dict(q) for o in options], | |
| "hidden_below_min_bits": below, | |
| "repos": repos, | |
| "geometry": detail.geometry.to_dict() if detail.geometry else None, | |
| "geometry_source": detail.geometry_source, | |
| "kv_cache_gb": ({c: round(detail.geometry.kv_bytes(c, KV_TYPES[req.kv_type]) / 1e9, 2) for c in CONTEXTS} | |
| if detail.geometry else None), | |
| "other_formats": other, | |
| "community": community.get(g.key), | |
| } | |
| async def events(request: Request): | |
| if not st.limiter.allow(client_key(request)): | |
| return JSONResponse({"error": "rate limited"}, status_code=429) | |
| raw = await request.body() | |
| if len(raw) > MAX_BODY: | |
| return JSONResponse({"error": "payload too large"}, status_code=413) | |
| try: | |
| body = json.loads(raw) | |
| except Exception: | |
| return JSONResponse({"error": "invalid json"}, status_code=400) | |
| try: | |
| batch = EventBatch.model_validate(body) | |
| except ValidationError as e: | |
| errs = [{"loc": [str(x) for x in err.get("loc", [])][-3:], "msg": str(err.get("msg"))[:160]} | |
| for err in e.errors()[:10]] | |
| return JSONResponse({"error": "invalid events", "details": errs}, status_code=422) | |
| if not st.collection_enabled: | |
| return {"accepted": 0} | |
| rows = [to_row(ev, st.catalogue, APP_VERSION, st.engine.version) for ev in batch.events] | |
| st.sink.add(rows) | |
| st.stats.invalidate() | |
| summary = [{"flags": r["suspicious_flags"], "generation_tps": r.get("generation_tps"), | |
| "prompt_tps": r.get("prompt_tps"), "backend": r.get("backend")} for r in rows] | |
| return {"accepted": len(rows), "rows": summary} | |
| def stats(): | |
| return st.stats.stats() | |
| app.mount("/static", StaticFiles(directory=STATIC), name="static") | |
| return app | |
| app = create_app() if os.environ.get("EXPLORER_NO_AUTOAPP") != "1" else None | |