Spaces:
Sleeping
Sleeping
File size: 2,452 Bytes
f779370 724a1ef 4235063 de526d2 f779370 de526d2 f779370 4235063 de526d2 d6c55c0 f779370 de526d2 | 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 | import os
import time
from functools import lru_cache
from typing import Any
import gradio as gr
import spaces
from fastapi import Header, HTTPException
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
from pydantic import BaseModel
MODEL_REPO = os.getenv("MODEL_REPO", "Qwen/Qwen2.5-1.5B-Instruct-GGUF")
MODEL_FILE = os.getenv("MODEL_FILE", "qwen2.5-1.5b-instruct-q4_k_m.gguf")
CHAT_FORMAT = os.getenv("CHAT_FORMAT", "chatml")
N_CTX = int(os.getenv("N_CTX", "2048"))
N_THREADS = int(os.getenv("N_THREADS", str(min(os.cpu_count() or 2, 2))))
SPACE_API_TOKEN = os.getenv("SPACE_API_TOKEN", "").strip()
app = gr.Server(title="Speaking Diff CPU LLM")
@spaces.GPU(duration=1)
@app.api()
def zerogpu_marker() -> str:
return "ok"
class ChatMessage(BaseModel):
role: str
content: str
class ChatCompletionRequest(BaseModel):
model: str | None = None
messages: list[ChatMessage]
temperature: float = 0.2
max_tokens: int = 500
@lru_cache(maxsize=1)
def get_llm() -> Llama:
model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
return Llama(
model_path=model_path,
n_ctx=N_CTX,
n_threads=N_THREADS,
n_gpu_layers=0,
chat_format=CHAT_FORMAT,
verbose=False,
)
def check_auth(authorization: str | None) -> None:
if not SPACE_API_TOKEN:
return
if authorization != f"Bearer {SPACE_API_TOKEN}":
raise HTTPException(status_code=401, detail="Unauthorized")
@app.get("/")
def root() -> dict[str, str]:
return {
"name": "Speaking Diff CPU LLM",
"model": f"{MODEL_REPO}/{MODEL_FILE}",
"chat_completions": "/chat/completions",
}
@app.get("/health")
def health() -> dict[str, Any]:
return {
"ok": True,
"model_repo": MODEL_REPO,
"model_file": MODEL_FILE,
"n_ctx": N_CTX,
"n_threads": N_THREADS,
}
@app.post("/chat/completions")
def chat_completions(request: ChatCompletionRequest, authorization: str | None = Header(default=None)) -> dict[str, Any]:
check_auth(authorization)
llm = get_llm()
result = llm.create_chat_completion(
messages=[message.model_dump() for message in request.messages],
temperature=request.temperature,
max_tokens=min(request.max_tokens, 700),
)
result["model"] = request.model or "local-qwen2.5-1.5b"
result.setdefault("created", int(time.time()))
return result
if __name__ == "__main__":
app.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))
|