File size: 5,740 Bytes
0990c3f | 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 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | from __future__ import annotations
import json
import os
import secrets
import threading
import time
import uuid
from contextlib import asynccontextmanager
from typing import Any, Literal
from fastapi import Depends, FastAPI, Header, HTTPException
from fastapi.concurrency import run_in_threadpool
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field
from tokenizers import Tokenizer
from nayara.artifacts import materialize_checkpoint
from nayara.generate import load_model, stream_generate
class CompletionRequest(BaseModel):
model: str = "nayara-v1"
prompt: str
max_tokens: int = Field(192, ge=1, le=512)
temperature: float = Field(0.2, ge=0, le=2)
top_k: int = Field(40, ge=1, le=32000)
stream: bool = False
class ChatMessage(BaseModel):
role: Literal["system", "user", "assistant"]
content: str
class ChatRequest(BaseModel):
model: str = "nayara-v1"
messages: list[ChatMessage]
max_tokens: int = Field(192, ge=1, le=512)
temperature: float = Field(0.2, ge=0, le=2)
top_k: int = Field(40, ge=1, le=32000)
stream: bool = False
runtime: dict[str, Any] = {"ready": False}
generation_lock = threading.Lock()
MODEL_ID = os.environ.get("NAYARA_MODEL_ID", "nayara-v1")
def _authorize(authorization: str | None = Header(default=None)) -> None:
expected = os.environ.get("NAYARA_API_KEY")
if not expected:
return
supplied = authorization.removeprefix("Bearer ") if authorization else ""
if not secrets.compare_digest(supplied, expected):
raise HTTPException(status_code=401, detail="invalid bearer token")
def _chat_prompt(messages: list[ChatMessage]) -> str:
if not messages or not any(message.role == "user" for message in messages):
raise HTTPException(status_code=422, detail="messages must contain at least one user message")
parts: list[str] = []
for message in messages:
if message.role == "system":
parts.append(f"<|user|>\nSystem guidance: {message.content.strip()}")
else:
parts.append(f"<|{message.role}|>\n{message.content.strip()}")
parts.append("<|assistant|>\n")
return "\n".join(parts)
def _generate(prompt: str, request: CompletionRequest | ChatRequest):
return stream_generate(
runtime["model"], runtime["tokenizer"], prompt,
max_new_tokens=request.max_tokens, temperature=request.temperature, top_k=request.top_k,
)
def _complete(prompt: str, request: CompletionRequest | ChatRequest) -> str:
final = ""
with generation_lock:
for final in _generate(prompt, request):
pass
return final
def _sse(prompt: str, request: CompletionRequest | ChatRequest, chat: bool):
request_id = f"nayara-{uuid.uuid4().hex}"
previous = ""
with generation_lock:
for partial in _generate(prompt, request):
delta, previous = partial[len(previous) :], partial
if chat:
payload = {"id": request_id, "object": "chat.completion.chunk", "model": MODEL_ID, "choices": [{"index": 0, "delta": {"content": delta}, "finish_reason": None}]}
else:
payload = {"id": request_id, "object": "text_completion", "model": MODEL_ID, "choices": [{"index": 0, "text": delta, "finish_reason": None}]}
yield f"data: {json.dumps(payload)}\n\n"
yield "data: [DONE]\n\n"
@asynccontextmanager
async def lifespan(_: FastAPI):
if os.environ.get("NAYARA_REQUIRE_API_KEY") == "1" and not os.environ.get("NAYARA_API_KEY"):
raise RuntimeError("NAYARA_REQUIRE_API_KEY=1 but NAYARA_API_KEY is unset")
checkpoint = await run_in_threadpool(materialize_checkpoint)
threads = int(os.environ.get("NAYARA_THREADS", max(1, (os.cpu_count() or 4) // 2)))
runtime["tokenizer"] = Tokenizer.from_file(str(checkpoint / "tokenizer.json"))
runtime["model"] = await run_in_threadpool(load_model, checkpoint, True, threads)
runtime.update({"ready": True, "checkpoint": str(checkpoint), "loaded_at": time.time()})
yield
runtime.clear()
runtime["ready"] = False
app = FastAPI(title="NAYARA API", version="1.0.0", lifespan=lifespan)
@app.get("/healthz")
def health() -> dict[str, str]:
return {"status": "ok"}
@app.get("/readyz")
def ready() -> dict[str, Any]:
if not runtime.get("ready"):
raise HTTPException(status_code=503, detail="model is not loaded")
return {"status": "ready", "model": MODEL_ID, "loaded_at": runtime["loaded_at"]}
@app.get("/v1/models", dependencies=[Depends(_authorize)])
def models() -> dict[str, Any]:
return {"object": "list", "data": [{"id": MODEL_ID, "object": "model", "owned_by": "nayara"}]}
@app.post("/v1/completions", dependencies=[Depends(_authorize)])
async def completions(request: CompletionRequest):
prompt = request.prompt
if request.stream:
return StreamingResponse(_sse(prompt, request, False), media_type="text/event-stream")
text = await run_in_threadpool(_complete, prompt, request)
return {"id": f"nayara-{uuid.uuid4().hex}", "object": "text_completion", "model": MODEL_ID, "choices": [{"index": 0, "text": text, "finish_reason": "stop"}]}
@app.post("/v1/chat/completions", dependencies=[Depends(_authorize)])
async def chat_completions(request: ChatRequest):
prompt = _chat_prompt(request.messages)
if request.stream:
return StreamingResponse(_sse(prompt, request, True), media_type="text/event-stream")
text = await run_in_threadpool(_complete, prompt, request)
return {"id": f"nayara-{uuid.uuid4().hex}", "object": "chat.completion", "model": MODEL_ID, "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}]}
|