Spaces:
Sleeping
Sleeping
File size: 3,597 Bytes
8b12a0a a847960 8b12a0a a847960 8b12a0a a847960 8b12a0a a847960 8b12a0a | 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 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | """
FastAPI backend that serves the chat model to the web frontend.
Discovers the available *_chat.pth checkpoints, loads each one lazily (and only
once), and exposes:
GET /api/models -> which chat models are available
POST /api/chat -> stream an assistant reply as Server-Sent Events
It also serves the built frontend (web/dist) when that directory exists.
"""
import json
from pathlib import Path
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from server.inference import ChatModel
ROOT = Path(__file__).resolve().parent.parent
SUFFIX = "_chat.pth"
def discover_checkpoints():
"""Map arch name -> checkpoint path for every *_chat.pth in the project root."""
found = {}
for path in sorted(ROOT.glob(f"*{SUFFIX}")):
arch = path.name[: -len(SUFFIX)]
if arch in ("gpt2", "gpt3", "gpt3.1"):
found[arch] = path
return found
_checkpoints = discover_checkpoints()
_models = {} # arch -> ChatModel, loaded on first use
def get_model(arch):
if arch not in _models:
_models[arch] = ChatModel("gpt3", str(_checkpoints[arch]))
return _models[arch]
app = FastAPI(title="GPT Chat")
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_methods=["*"],
allow_headers=["*"],
)
class Message(BaseModel):
role: str
content: str
class ChatRequest(BaseModel):
messages: list[Message]
model: str | None = None
max_new_tokens: int = 200
top_k: int = 50
@app.get("/api/models")
def list_models():
return {"models": list(_checkpoints.keys())}
# Headers that keep the SSE response unbuffered through Hugging Face's proxy, so
# tokens reach the browser as they are generated and a client disconnect is
# detected promptly (which stops generation and frees the per-model lock instead
# of leaving an abandoned request running and blocking everyone else).
SSE_HEADERS = {
"Cache-Control": "no-cache",
"X-Accel-Buffering": "no",
"Connection": "keep-alive",
}
@app.post("/api/chat")
def chat(req: ChatRequest):
arch = req.model or (next(iter(_checkpoints), None))
def sse(obj):
return f"data: {json.dumps(obj)}\n\n"
if arch not in _checkpoints:
return StreamingResponse(
iter([sse({"error": f"unknown model: {arch}"})]),
media_type="text/event-stream",
headers=SSE_HEADERS,
)
model = get_model(arch)
messages = [m.model_dump() for m in req.messages]
max_new_tokens = max(1, min(req.max_new_tokens, 1000))
top_k = max(1, min(req.top_k, 200))
def event_stream():
# Flush an SSE comment up front so the proxy starts forwarding the stream
# immediately rather than buffering it until generation finishes.
yield ": connected\n\n"
try:
for delta in model.generate_stream(messages, max_new_tokens, top_k):
yield sse({"delta": delta})
yield sse({"done": True})
except Exception as e: # surface generation errors to the client
yield sse({"error": str(e)})
return StreamingResponse(
event_stream(), media_type="text/event-stream", headers=SSE_HEADERS
)
# Serve the built frontend if present (production single-origin deploy).
# Mounted last so it does not shadow the /api routes above.
_dist = ROOT / "web" / "dist"
if _dist.is_dir():
app.mount("/", StaticFiles(directory=str(_dist), html=True), name="static")
|