| """The unified local server: decisions, generation, and decide-then-generate from one cache. |
| |
| uv run --group serve python -m serve.server --artifacts /path/to/artifacts --port 8008 |
| |
| POST /v1/systemone typed decisions with probabilities and the router's flag |
| POST /v1/chat/completions base-model generation, adapter off |
| POST /v1/auto one question: decide, and decode from the state cache if flagged |
| |
| Plain PyTorch, one request on the model at a time: correctness before speed. |
| """ |
|
|
| import argparse |
| import time |
| import uuid |
| from pathlib import Path |
|
|
| from fastapi import FastAPI, HTTPException |
| from fastapi.responses import HTMLResponse |
| from pydantic import BaseModel, Field |
|
|
| from dmm.engine import Engine |
|
|
|
|
| class SystemOneRequest(BaseModel): |
| state: str |
| questions: dict[str, dict] = Field(min_length=1) |
| model: str = "dev-4b" |
|
|
|
|
| class ChatRequest(BaseModel): |
| messages: list[dict] |
| max_tokens: int = 512 |
| model: str = "dev-4b" |
|
|
|
|
| class AutoRequest(BaseModel): |
| state: str |
| question: dict |
| name: str = "q" |
| max_tokens: int = 512 |
|
|
|
|
| def create_app(engine: Engine) -> FastAPI: |
| app = FastAPI(title="Dev-4B") |
|
|
| @app.get("/", response_class=HTMLResponse) |
| def playground(): |
| """A small browser page for trying the three endpoints by hand.""" |
| return (Path(__file__).parent / "playground.html").read_text() |
|
|
| @app.post("/v1/systemone") |
| def systemone(req: SystemOneRequest): |
| t0 = time.perf_counter() |
| try: |
| answers = engine.decide(req.state, req.questions) |
| except ValueError as e: |
| raise HTTPException(422, str(e)) from e |
| return { |
| "model": req.model, |
| "answers": answers, |
| "latency_ms": round((time.perf_counter() - t0) * 1000, 1), |
| } |
|
|
| @app.post("/v1/chat/completions") |
| def chat(req: ChatRequest): |
| tokens = engine.chat(req.messages, req.max_tokens) |
| text = engine.tok.decode(tokens, skip_special_tokens=True) |
| return { |
| "id": f"chatcmpl-{uuid.uuid4().hex[:12]}", |
| "object": "chat.completion", |
| "model": req.model, |
| "tokens": tokens, |
| "choices": [ |
| { |
| "index": 0, |
| "message": {"role": "assistant", "content": text}, |
| "finish_reason": "stop", |
| } |
| ], |
| "usage": {"completion_tokens": len(tokens)}, |
| } |
|
|
| @app.post("/v1/auto") |
| def auto(req: AutoRequest): |
| try: |
| return engine.auto(req.state, req.name, req.question, req.max_tokens) |
| except ValueError as e: |
| raise HTTPException(422, str(e)) from e |
|
|
| return app |
|
|
|
|
| def main(): |
| import uvicorn |
| import yaml |
|
|
| from dmm.load import load |
|
|
| ap = argparse.ArgumentParser() |
| ap.add_argument("--config", default="configs/stage6.yaml") |
| ap.add_argument("--artifacts", required=True, help="dir with best.pt and temperatures.json") |
| ap.add_argument("--router", default=None, help="dir with router.pt and operating_point.json") |
| ap.add_argument("--device", default="cuda") |
| ap.add_argument("--host", default="127.0.0.1") |
| ap.add_argument("--port", type=int, default=8008) |
| a = ap.parse_args() |
| cfg = yaml.safe_load(open(a.config)) |
| base = yaml.safe_load(open(cfg["base_config"])) |
| model, tok = load(base["model"], device=a.device) |
| engine = Engine.from_artifacts(model, tok, cfg, a.artifacts, a.router or a.artifacts) |
| uvicorn.run(create_app(engine), host=a.host, port=a.port) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
