Spaces:
Paused
Paused
| """booper-chat, served. | |
| A 34M-parameter model answers in well under a second on CPU, so the whole | |
| design here is about not squandering that: weights are baked into the image | |
| (no cold download), the page is one static file with no framework to boot, | |
| and tokens go out over SSE as they are produced rather than after the reply | |
| is finished. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import os | |
| import threading | |
| from pathlib import Path | |
| from fastapi import FastAPI | |
| from fastapi.responses import FileResponse, JSONResponse, StreamingResponse | |
| from fastapi.staticfiles import StaticFiles | |
| from pydantic import BaseModel, Field | |
| from booper.engine import Booper, SamplingParams, Turn | |
| ROOT = Path(__file__).parent | |
| MODEL_DIR = Path(os.environ.get("BOOPER_MODEL_DIR", ROOT / "weights")) | |
| MODEL_ID = "ProCreations/booper-chat" | |
| # One generation at a time. The Space runs on 2 vCPUs and torch already uses | |
| # both for a single decode, so overlapping two requests makes each of them | |
| # slower rather than making the pair faster; queueing keeps every reply fast | |
| # and the tok/s number honest. | |
| _generation_lock = threading.Lock() | |
| MAX_QUEUE = int(os.environ.get("BOOPER_MAX_QUEUE", "24")) | |
| _waiting = threading.Semaphore(MAX_QUEUE) | |
| # Input caps. The 1024-token window truncates anything longer anyway, but the | |
| # tokenizer would still have to chew through whatever arrives first. | |
| MAX_TURNS = 400 | |
| MAX_CHARS = 60_000 | |
| booper = Booper(MODEL_DIR / "model.pt", MODEL_DIR / "tokenizer.json") | |
| print( | |
| f"[booper] {booper.n_params:,} params · block {booper.block_size} · " | |
| f"vocab {booper.tok.vocab_size} · backend {booper.backend}", | |
| flush=True, | |
| ) | |
| app = FastAPI(title="booper-chat", docs_url=None, redoc_url=None) | |
| class Message(BaseModel): | |
| role: str = "user" | |
| content: str = "" | |
| class ChatRequest(BaseModel): | |
| messages: list[Message] = Field(default_factory=list) | |
| max_new_tokens: int = 96 | |
| temperature: float = 0.7 | |
| top_k: int = 40 | |
| top_p: float = 0.9 | |
| repetition_penalty: float = 1.15 | |
| frequency_penalty: float = 0.12 | |
| presence_penalty: float = 0.0 | |
| no_repeat_ngram_size: int = 3 | |
| seed: int | None = None | |
| def _sse(event: str, payload: dict) -> str: | |
| return f"event: {event}\ndata: {json.dumps(payload, separators=(',', ':'))}\n\n" | |
| def meta() -> JSONResponse: | |
| return JSONResponse( | |
| { | |
| "model": MODEL_ID, | |
| "params": booper.n_params, | |
| "context": booper.block_size, | |
| "vocab": booper.tok.vocab_size, | |
| "layers": booper.config.n_layer, | |
| "width": booper.config.n_embd, | |
| "heads": booper.config.n_head, | |
| "quantized": booper.quantized, | |
| "defaults": SamplingParams().__dict__, | |
| } | |
| ) | |
| def chat(request: ChatRequest) -> StreamingResponse: | |
| turns = [ | |
| Turn("assistant" if m.role == "assistant" else "user", m.content) | |
| for m in request.messages[-MAX_TURNS:] | |
| if m.content and m.content.strip() | |
| ] | |
| total = sum(len(t.content) for t in turns) | |
| while total > MAX_CHARS and len(turns) > 1: | |
| total -= len(turns.pop(0).content) | |
| if turns: | |
| turns[-1] = Turn(turns[-1].role, turns[-1].content[:MAX_CHARS]) | |
| params = SamplingParams( | |
| max_new_tokens=request.max_new_tokens, | |
| temperature=request.temperature, | |
| top_k=request.top_k, | |
| top_p=request.top_p, | |
| repetition_penalty=request.repetition_penalty, | |
| frequency_penalty=request.frequency_penalty, | |
| presence_penalty=request.presence_penalty, | |
| no_repeat_ngram_size=request.no_repeat_ngram_size, | |
| seed=request.seed, | |
| ) | |
| def body(): | |
| if not turns: | |
| yield _sse("error", {"message": "say something first"}) | |
| return | |
| if not _waiting.acquire(blocking=False): | |
| yield _sse("error", {"message": "too many people talking to booper at once — try again"}) | |
| return | |
| try: | |
| queued = not _generation_lock.acquire(blocking=False) | |
| if queued: | |
| yield _sse("queued", {}) | |
| _generation_lock.acquire() | |
| try: | |
| for kind, payload in booper.stream(turns, params): | |
| if kind == "token": | |
| yield _sse("token", {"t": payload}) | |
| else: | |
| yield _sse("done", payload) | |
| finally: | |
| _generation_lock.release() | |
| except Exception as exc: # never leave the client hanging on a dead stream | |
| yield _sse("error", {"message": f"{type(exc).__name__}: {exc}"}) | |
| finally: | |
| _waiting.release() | |
| return StreamingResponse( | |
| body(), | |
| media_type="text/event-stream", | |
| headers={ | |
| "Cache-Control": "no-cache, no-transform", | |
| "Connection": "keep-alive", | |
| # Spaces sit behind a proxy that will happily buffer a whole SSE | |
| # stream and deliver it as one lump, which would silently undo | |
| # every bit of the streaming above. | |
| "X-Accel-Buffering": "no", | |
| }, | |
| ) | |
| def index() -> FileResponse: | |
| return FileResponse(ROOT / "static" / "index.html", headers={"Cache-Control": "no-cache"}) | |
| app.mount("/", StaticFiles(directory=ROOT / "static"), name="static") | |