Spaces:
Running
Running
| """Run: uv run python server.py --port 8043.""" | |
| from __future__ import annotations | |
| import argparse | |
| import threading | |
| from contextlib import asynccontextmanager | |
| from pathlib import Path | |
| from typing import Literal | |
| import uvicorn | |
| from fastapi import FastAPI, HTTPException, Query | |
| from fastapi.responses import FileResponse, Response | |
| from fastapi.staticfiles import StaticFiles | |
| from pydantic import BaseModel, ConfigDict, Field | |
| from runtime import Runtime, MAX_TOKENS, MODEL_ID, GEMMA_ID, MODELS | |
| ROOT = Path(__file__).parent | |
| ModelChoice = Literal[MODEL_ID, GEMMA_ID] | |
| runtimes = {model_id: Runtime(model_id) for model_id in MODELS} | |
| class Message(BaseModel): | |
| role: Literal["system", "user", "assistant"] | |
| content: str = Field(max_length=50000) | |
| class Input(BaseModel): | |
| model_config = ConfigDict(extra="forbid") | |
| model: ModelChoice = MODEL_ID | |
| mode: Literal["chat", "raw"] = "chat" | |
| messages: list[Message] = Field(default_factory=list, max_length=40) | |
| text: str = Field(default="", max_length=100000) | |
| thinking: bool = False | |
| max_new_tokens: int = Field(default=0, ge=0, le=128) | |
| class AppendInput(BaseModel): | |
| model_config = ConfigDict(extra="forbid") | |
| text: str | None = Field(default=None, max_length=20000) | |
| async def lifespan(app): | |
| def load_models(): | |
| for runtime in runtimes.values(): | |
| runtime.load() | |
| threading.Thread(target=load_models, daemon=True).start() | |
| yield | |
| app = FastAPI(lifespan=lifespan) | |
| def status(model: ModelChoice = MODEL_ID): | |
| return dict(**runtimes[model].status(), models=[dict(id=mid, **spec) for mid, spec in MODELS.items()]) | |
| def prepare(payload: Input): | |
| runtime = runtimes[payload.model] | |
| try: | |
| text, ids = runtime.prepare(payload) | |
| return dict(text=text, count=len(ids), tokens=[runtime.token(i) for i in ids], | |
| max_tokens=MAX_TOKENS, fits=len(ids) + payload.max_new_tokens <= MAX_TOKENS) | |
| except ValueError as exc: | |
| raise HTTPException(422, str(exc)) from exc | |
| def analyze(payload: Input): | |
| try: | |
| return runtimes[payload.model].analyze(payload) | |
| except BlockingIOError as exc: | |
| raise HTTPException(409, str(exc)) from exc | |
| except ValueError as exc: | |
| raise HTTPException(422, str(exc)) from exc | |
| except Exception as exc: | |
| import traceback | |
| traceback.print_exc() | |
| raise HTTPException(500, f"Computation failed: {exc}") from exc | |
| def get_session(sid): | |
| for runtime in runtimes.values(): | |
| try: | |
| return runtime, runtime.session(sid) | |
| except KeyError: | |
| continue | |
| raise HTTPException(404, "This result has been evicted from memory. Run the analysis again.") | |
| def inspect(sid: str, layer: int = Query(ge=0), query: int = Query(ge=0)): | |
| runtime, s = get_session(sid) | |
| if layer >= s["response"]["layers"]: | |
| raise HTTPException(422, "Layer is outside this model.") | |
| if query >= len(s["response"]["tokens"]): | |
| raise HTTPException(422, "Token position is outside the context.") | |
| return runtime.inspect(sid, layer, query) | |
| def append_preview(sid: str, payload: AppendInput): | |
| runtime, session = get_session(sid) | |
| try: | |
| ids = runtime.append_ids(sid, payload.text) | |
| except ValueError as exc: | |
| raise HTTPException(422, str(exc)) from exc | |
| total = len(session["response"]["tokens"]) + len(ids) | |
| return dict(tokens=[runtime.token(tid) for tid in ids], count=len(ids), total=total, | |
| max_tokens=MAX_TOKENS, fits=total <= MAX_TOKENS) | |
| def continue_run(sid: str, payload: AppendInput | None = None): | |
| runtime, s = get_session(sid) | |
| previous = s["response"] | |
| ids = [token["id"] for token in previous["tokens"]] | |
| try: | |
| # Preserve the exact token IDs, including a possible partial UTF-8 token. | |
| text = payload.text if payload is not None else None | |
| suffix = runtime.append_ids(sid, text) | |
| sources = previous["token_sources"] + ["predicted" if text is None else "manual"] * len(suffix) | |
| return runtime.analyze(Input(model=runtime.model_id, mode="raw"), | |
| prefix=(previous["context"], ids + suffix), | |
| input_count=previous["input_count"], token_sources=sources) | |
| except BlockingIOError as exc: | |
| raise HTTPException(409, str(exc)) from exc | |
| except ValueError as exc: | |
| raise HTTPException(422, str(exc)) from exc | |
| def attention(sid: str, layer: int = Query(ge=0), head: int | Literal["mean"] = Query(...)): | |
| _, s = get_session(sid) | |
| if layer >= s["response"]["layers"] or (head != "mean" and not 0 <= head < s["response"]["heads"]): | |
| raise HTTPException(422, "Layer or head is outside this model.") | |
| heads = s["attention"][layer] | |
| # Average stored post-softmax probabilities, accumulating in float32. | |
| matrix = heads.mean(axis=0, dtype="<f4") if head == "mean" else heads[head].astype("<f4") | |
| return Response(matrix.tobytes(), media_type="application/octet-stream", | |
| headers={"X-Token-Count": str(matrix.shape[0]), "Cache-Control": "private, max-age=3600"}) | |
| def attention_layers(sid: str, query: int = Query(ge=0), head: int | Literal["mean"] = "mean"): | |
| runtime, session = get_session(sid) | |
| if query >= len(session["response"]["tokens"]): | |
| raise HTTPException(422, "Token position is outside the context.") | |
| if head != "mean" and not 0 <= head < session["response"]["heads"]: | |
| raise HTTPException(422, "Head is outside this model.") | |
| return runtime.attention_layers(sid, query, head) | |
| def index(): | |
| return FileResponse(ROOT / "static" / "index.html") | |
| app.mount("/", StaticFiles(directory=ROOT / "static"), name="static") | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--port", type=int, default=8043) | |
| parser.add_argument("--host", default="127.0.0.1") | |
| args = parser.parse_args() | |
| uvicorn.run(app, host=args.host, port=args.port) | |