Spaces:
Sleeping
Sleeping
File size: 2,809 Bytes
a7f7a81 d240987 a7f7a81 d240987 a7f7a81 d240987 a7f7a81 d240987 a7f7a81 d240987 a7f7a81 | 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 | #!/usr/bin/env python3
import os, sys, time, threading
from pathlib import Path
from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Optional
import uvicorn
MODEL_PATH = Path(os.getenv("HOME", "/home/user")) / "models" / os.getenv("MODEL_FILE", "qwen2.5-7b-instruct-q3_k_m.gguf")
MODEL_NAME = os.getenv("SERVED_MODEL_NAME", "qwen")
API_KEY = os.getenv("API_KEY", "")
app = FastAPI(title="Qwen 2.5 API")
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True,
allow_methods=["*"], allow_headers=["*"])
llm = None
llm_lock = threading.Lock()
def get_llm():
global llm
if llm is not None:
return llm
with llm_lock:
if llm is not None:
return llm
print(f"Loading model from {MODEL_PATH}...", flush=True)
t0 = time.time()
from llama_cpp import Llama
llm = Llama(
model_path=str(MODEL_PATH),
n_gpu_layers=0,
n_ctx=int(os.getenv("N_CTX", "8192")),
n_threads=int(os.getenv("N_THREADS", "2")),
n_batch=int(os.getenv("N_BATCH", "256")),
n_ubatch=int(os.getenv("N_UBATCH", "128")),
verbose=False,
)
print(f"Model loaded in {time.time()-t0:.1f}s", flush=True)
return llm
class Message(BaseModel):
role: str; content: str
class ChatRequest(BaseModel):
model: str = MODEL_NAME
messages: List[Message]
max_tokens: Optional[int] = 512
temperature: Optional[float] = 0.7
async def check_auth(request: Request):
if API_KEY:
auth = request.headers.get("Authorization", "")
if auth != f"Bearer {API_KEY}":
raise HTTPException(status_code=401, detail="Invalid API key")
@app.get("/")
@app.get("/v1/models")
async def models():
return {"data": [{"id": MODEL_NAME, "object": "model", "owned_by": "local"}]}
@app.get("/health")
async def health():
return {"status": "ok"}
@app.post("/v1/chat/completions")
async def chat(req: ChatRequest, request: Request):
await check_auth(request)
instance = get_llm()
messages = [{"role": m.role, "content": m.content} for m in req.messages]
r = instance.create_chat_completion(
messages=messages, max_tokens=req.max_tokens,
temperature=req.temperature,
)
return {
"id": "chatcmpl-1", "object": "chat.completion",
"model": req.model,
"choices": [{"index": 0, "message": r["choices"][0]["message"],
"finish_reason": r["choices"][0]["finish_reason"]}]
}
if __name__ == "__main__":
print("Starting uvicorn (lazy model load)...", flush=True)
uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info")
|