Spaces:
Running on Zero
Running on Zero
| 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") | |
| 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 | |
| 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") | |
| def root() -> dict[str, str]: | |
| return { | |
| "name": "Speaking Diff CPU LLM", | |
| "model": f"{MODEL_REPO}/{MODEL_FILE}", | |
| "chat_completions": "/chat/completions", | |
| } | |
| def health() -> dict[str, Any]: | |
| return { | |
| "ok": True, | |
| "model_repo": MODEL_REPO, | |
| "model_file": MODEL_FILE, | |
| "n_ctx": N_CTX, | |
| "n_threads": N_THREADS, | |
| } | |
| 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"))) | |