Spaces:
Running
Running
| import os | |
| import json | |
| import time | |
| import uuid | |
| import asyncio | |
| from typing import Any | |
| from fastapi import FastAPI, Request, HTTPException, Depends | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from fastapi.responses import JSONResponse, StreamingResponse | |
| from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials | |
| from huggingface_hub import hf_hub_download | |
| from llama_cpp import Llama | |
| # ============================================================ | |
| # CONFIG | |
| # ============================================================ | |
| APP_NAME = "Varsh LLM" | |
| MODEL_ID = "qwen2.5-coder-7b-instruct-q4_k_m" | |
| MODEL_REPO = "Qwen/Qwen2.5-Coder-7B-Instruct-GGUF" | |
| MODEL_FILE = "qwen2.5-coder-7b-instruct-q4_k_m.gguf" | |
| HOST = "0.0.0.0" | |
| PORT = int(os.getenv("PORT", "7860")) | |
| CUSTOM_API_KEY = os.getenv("CUSTOM_API_KEY") | |
| if not CUSTOM_API_KEY: | |
| raise RuntimeError( | |
| "CUSTOM_API_KEY is not configured. " | |
| "Add CUSTOM_API_KEY in Hugging Face Space Settings > Secrets." | |
| ) | |
| # ============================================================ | |
| # CONTEXT SETTINGS | |
| # ============================================================ | |
| # 16 GB RAM + 7B Q4 CPU deployment. | |
| # 8192 is a safer target than 16384. | |
| CONTEXT_SIZE = 8192 | |
| # Reserve room for the model's answer. | |
| RESERVED_OUTPUT_TOKENS = 1024 | |
| # Maximum amount of input context we allow. | |
| MAX_INPUT_TOKENS = ( | |
| CONTEXT_SIZE - RESERVED_OUTPUT_TOKENS | |
| ) | |
| # ============================================================ | |
| # FASTAPI | |
| # ============================================================ | |
| app = FastAPI( | |
| title=APP_NAME, | |
| version="4.0.0", | |
| ) | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_credentials=True, | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| security = HTTPBearer(auto_error=False) | |
| # ============================================================ | |
| # AUTH | |
| # ============================================================ | |
| def verify_api_key( | |
| credentials: HTTPAuthorizationCredentials = | |
| Depends(security), | |
| ): | |
| if credentials is None: | |
| raise HTTPException( | |
| status_code=401, | |
| detail="Missing Authorization header.", | |
| ) | |
| if credentials.scheme.lower() != "bearer": | |
| raise HTTPException( | |
| status_code=401, | |
| detail="Authorization must use Bearer authentication.", | |
| ) | |
| if credentials.credentials != CUSTOM_API_KEY: | |
| raise HTTPException( | |
| status_code=401, | |
| detail="Invalid API key.", | |
| ) | |
| return True | |
| # ============================================================ | |
| # MODEL | |
| # ============================================================ | |
| print("=" * 70) | |
| print("Starting Varsh LLM") | |
| print("=" * 70) | |
| print(f"Model repository: {MODEL_REPO}") | |
| print(f"Model file: {MODEL_FILE}") | |
| print(f"Context size: {CONTEXT_SIZE}") | |
| print(f"Max input: {MAX_INPUT_TOKENS}") | |
| print(f"Reserved output: {RESERVED_OUTPUT_TOKENS}") | |
| print("Checking/downloading model...") | |
| model_path = hf_hub_download( | |
| repo_id=MODEL_REPO, | |
| filename=MODEL_FILE, | |
| ) | |
| print(f"Model path: {model_path}") | |
| print("Loading Qwen 2.5 Coder 7B...") | |
| llm = Llama( | |
| model_path=model_path, | |
| # CPU deployment | |
| n_ctx=CONTEXT_SIZE, | |
| n_threads=2, | |
| n_threads_batch=2, | |
| # Moderate batch for 2 vCPU. | |
| n_batch=256, | |
| use_mmap=True, | |
| use_mlock=False, | |
| verbose=False, | |
| ) | |
| print("Model loaded successfully.") | |
| print("=" * 70) | |
| print("VARSH LLM API READY") | |
| print("=" * 70) | |
| print(f"Listening on: {HOST}:{PORT}") | |
| print("Root endpoint: /") | |
| print("Models endpoint: /v1/models") | |
| print("Chat endpoint: /v1/chat/completions") | |
| print("Authentication: Bearer API key") | |
| print("Streaming: enabled") | |
| print(f"Context: {CONTEXT_SIZE}") | |
| print("=" * 70) | |
| # llama.cpp generation is serialized because this is | |
| # a 2-vCPU CPU deployment. | |
| generation_lock = asyncio.Lock() | |
| # ============================================================ | |
| # CONTENT NORMALIZATION | |
| # ============================================================ | |
| def content_to_text(content: Any) -> str: | |
| if content is None: | |
| return "" | |
| if isinstance(content, str): | |
| return content | |
| if isinstance(content, list): | |
| parts = [] | |
| for item in content: | |
| if isinstance(item, str): | |
| parts.append(item) | |
| continue | |
| if not isinstance(item, dict): | |
| continue | |
| item_type = item.get("type") | |
| if item_type == "text": | |
| text = item.get("text", "") | |
| if isinstance(text, str): | |
| parts.append(text) | |
| continue | |
| # Some clients use: | |
| # {"text": "..."} | |
| if isinstance( | |
| item.get("text"), | |
| str, | |
| ): | |
| parts.append( | |
| item["text"] | |
| ) | |
| return "\n".join(parts) | |
| if isinstance(content, dict): | |
| if isinstance( | |
| content.get("text"), | |
| str, | |
| ): | |
| return content["text"] | |
| if isinstance( | |
| content.get("content"), | |
| str, | |
| ): | |
| return content["content"] | |
| return json.dumps( | |
| content, | |
| ensure_ascii=False, | |
| ) | |
| return str(content) | |
| # ============================================================ | |
| # MESSAGE NORMALIZATION | |
| # ============================================================ | |
| def normalize_messages(messages: Any) -> list: | |
| if not isinstance(messages, list): | |
| raise HTTPException( | |
| status_code=400, | |
| detail="'messages' must be an array.", | |
| ) | |
| result = [] | |
| for message in messages: | |
| if not isinstance(message, dict): | |
| continue | |
| role = message.get( | |
| "role", | |
| "user", | |
| ) | |
| if not isinstance(role, str): | |
| role = "user" | |
| content = content_to_text( | |
| message.get( | |
| "content", | |
| "", | |
| ) | |
| ) | |
| normalized = { | |
| "role": role, | |
| "content": content, | |
| } | |
| if isinstance( | |
| message.get("name"), | |
| str, | |
| ): | |
| normalized["name"] = ( | |
| message["name"] | |
| ) | |
| if isinstance( | |
| message.get("tool_call_id"), | |
| str, | |
| ): | |
| normalized["tool_call_id"] = ( | |
| message["tool_call_id"] | |
| ) | |
| result.append(normalized) | |
| if not result: | |
| raise HTTPException( | |
| status_code=400, | |
| detail="'messages' cannot be empty.", | |
| ) | |
| return result | |
| # ============================================================ | |
| # TOKEN COUNTING | |
| # ============================================================ | |
| def count_message_tokens(messages: list) -> int: | |
| """ | |
| Use llama.cpp's actual tokenizer instead of estimating | |
| tokens using character length. | |
| """ | |
| try: | |
| # llama.cpp's chat template will ultimately create | |
| # a prompt. We use a conservative approximation here. | |
| total = 0 | |
| for message in messages: | |
| role = message.get( | |
| "role", | |
| "", | |
| ) | |
| content = message.get( | |
| "content", | |
| "", | |
| ) | |
| text = ( | |
| f"<|im_start|>{role}\n" | |
| f"{content}" | |
| f"<|im_end|>\n" | |
| ) | |
| tokens = llm.tokenize( | |
| text.encode( | |
| "utf-8" | |
| ), | |
| add_bos=False, | |
| ) | |
| total += len(tokens) | |
| # Assistant generation marker. | |
| assistant_tokens = llm.tokenize( | |
| b"<|im_start|>assistant\n", | |
| add_bos=False, | |
| ) | |
| total += len( | |
| assistant_tokens | |
| ) | |
| return total | |
| except Exception: | |
| # Safe fallback. | |
| characters = sum( | |
| len( | |
| str( | |
| m.get( | |
| "content", | |
| "", | |
| ) | |
| ) | |
| ) | |
| for m in messages | |
| ) | |
| return max( | |
| 1, | |
| characters // 4, | |
| ) | |
| # ============================================================ | |
| # CONTEXT MANAGEMENT | |
| # ============================================================ | |
| def fit_messages_to_context( | |
| messages: list, | |
| max_tokens: int, | |
| ) -> tuple[list, int]: | |
| if not messages: | |
| return messages, 0 | |
| current_tokens = count_message_tokens( | |
| messages | |
| ) | |
| if current_tokens <= max_tokens: | |
| return messages, current_tokens | |
| print( | |
| f"Context too large: " | |
| f"{current_tokens} tokens > " | |
| f"{max_tokens}. " | |
| f"Trimming old messages." | |
| ) | |
| # -------------------------------------------------------- | |
| # Preserve system messages. | |
| # -------------------------------------------------------- | |
| system_messages = [ | |
| m | |
| for m in messages | |
| if m.get("role") == "system" | |
| ] | |
| non_system_messages = [ | |
| m | |
| for m in messages | |
| if m.get("role") != "system" | |
| ] | |
| # -------------------------------------------------------- | |
| # Keep newest messages first. | |
| # -------------------------------------------------------- | |
| kept = list(system_messages) | |
| # Add newest conversation messages backwards. | |
| selected = [] | |
| for message in reversed( | |
| non_system_messages | |
| ): | |
| candidate = ( | |
| system_messages | |
| + list(reversed( | |
| selected + [message] | |
| )) | |
| ) | |
| tokens = count_message_tokens( | |
| candidate | |
| ) | |
| if tokens <= max_tokens: | |
| selected.append(message) | |
| else: | |
| break | |
| selected.reverse() | |
| final_messages = ( | |
| system_messages | |
| + selected | |
| ) | |
| final_tokens = count_message_tokens( | |
| final_messages | |
| ) | |
| # -------------------------------------------------------- | |
| # If system prompt itself is too large, | |
| # truncate the system content. | |
| # -------------------------------------------------------- | |
| if final_tokens > max_tokens: | |
| print( | |
| "System context is too large. " | |
| "Truncating system message." | |
| ) | |
| new_messages = [] | |
| remaining = max_tokens | |
| for message in final_messages: | |
| if ( | |
| message.get("role") | |
| != "system" | |
| ): | |
| continue | |
| content = message.get( | |
| "content", | |
| "", | |
| ) | |
| # Tokenize system content. | |
| tokens = llm.tokenize( | |
| content.encode( | |
| "utf-8" | |
| ), | |
| add_bos=False, | |
| ) | |
| if len(tokens) <= remaining: | |
| new_messages.append( | |
| message | |
| ) | |
| remaining -= len(tokens) | |
| else: | |
| # Approximate character cutoff. | |
| ratio = ( | |
| remaining | |
| / max( | |
| len(tokens), | |
| 1, | |
| ) | |
| ) | |
| cutoff = max( | |
| 1, | |
| int( | |
| len(content) | |
| * ratio | |
| ), | |
| ) | |
| truncated = content[ | |
| :cutoff | |
| ] | |
| new_messages.append( | |
| { | |
| "role": "system", | |
| "content": truncated, | |
| } | |
| ) | |
| remaining = 0 | |
| break | |
| # Add as many newest messages as fit. | |
| for message in selected: | |
| candidate = ( | |
| new_messages | |
| + [message] | |
| ) | |
| if ( | |
| count_message_tokens( | |
| candidate | |
| ) | |
| <= max_tokens | |
| ): | |
| new_messages.append( | |
| message | |
| ) | |
| final_messages = new_messages | |
| final_tokens = count_message_tokens( | |
| final_messages | |
| ) | |
| print( | |
| f"Context after trimming: " | |
| f"{final_tokens} tokens" | |
| ) | |
| return ( | |
| final_messages, | |
| final_tokens, | |
| ) | |
| # ============================================================ | |
| # HELPERS | |
| # ============================================================ | |
| def make_id(): | |
| return f"chatcmpl-{uuid.uuid4().hex}" | |
| def safe_int( | |
| value, | |
| default, | |
| ): | |
| try: | |
| return int(value) | |
| except Exception: | |
| return default | |
| def safe_float( | |
| value, | |
| default, | |
| ): | |
| try: | |
| return float(value) | |
| except Exception: | |
| return default | |
| # ============================================================ | |
| # ROOT | |
| # ============================================================ | |
| async def root(): | |
| return { | |
| "status": "healthy", | |
| "service": APP_NAME, | |
| "security": "enabled", | |
| "model": MODEL_ID, | |
| "endpoint": "/v1/chat/completions", | |
| "models_endpoint": "/v1/models", | |
| "streaming": True, | |
| "context_size": CONTEXT_SIZE, | |
| } | |
| # ============================================================ | |
| # HEALTH | |
| # ============================================================ | |
| async def health(): | |
| return { | |
| "status": "ok", | |
| "model": MODEL_ID, | |
| "context_size": CONTEXT_SIZE, | |
| } | |
| # ============================================================ | |
| # MODELS | |
| # ============================================================ | |
| async def models(): | |
| return { | |
| "object": "list", | |
| "data": [ | |
| { | |
| "id": MODEL_ID, | |
| "object": "model", | |
| "created": int(time.time()), | |
| "owned_by": "varsshancoder", | |
| } | |
| ], | |
| } | |
| # ============================================================ | |
| # CHAT | |
| # ============================================================ | |
| async def chat_completions( | |
| request: Request, | |
| ): | |
| # -------------------------------------------------------- | |
| # JSON | |
| # -------------------------------------------------------- | |
| try: | |
| body = await request.json() | |
| except Exception: | |
| raise HTTPException( | |
| status_code=400, | |
| detail="Invalid JSON request body.", | |
| ) | |
| if not isinstance(body, dict): | |
| raise HTTPException( | |
| status_code=400, | |
| detail="Request body must be a JSON object.", | |
| ) | |
| # -------------------------------------------------------- | |
| # Messages | |
| # -------------------------------------------------------- | |
| messages = normalize_messages( | |
| body.get( | |
| "messages", | |
| [], | |
| ) | |
| ) | |
| # -------------------------------------------------------- | |
| # Model | |
| # -------------------------------------------------------- | |
| requested_model = body.get( | |
| "model", | |
| MODEL_ID, | |
| ) | |
| if requested_model == f"varsh/{MODEL_ID}": | |
| requested_model = MODEL_ID | |
| # -------------------------------------------------------- | |
| # Parameters | |
| # -------------------------------------------------------- | |
| temperature = safe_float( | |
| body.get( | |
| "temperature", | |
| 0.2, | |
| ), | |
| 0.2, | |
| ) | |
| top_p = safe_float( | |
| body.get( | |
| "top_p", | |
| 0.95, | |
| ), | |
| 0.95, | |
| ) | |
| max_tokens = safe_int( | |
| body.get( | |
| "max_tokens", | |
| 512, | |
| ), | |
| 512, | |
| ) | |
| stream = bool( | |
| body.get( | |
| "stream", | |
| False, | |
| ) | |
| ) | |
| temperature = max( | |
| 0.0, | |
| min( | |
| temperature, | |
| 2.0, | |
| ), | |
| ) | |
| top_p = max( | |
| 0.01, | |
| min( | |
| top_p, | |
| 1.0, | |
| ), | |
| ) | |
| # -------------------------------------------------------- | |
| # Context management | |
| # -------------------------------------------------------- | |
| messages, prompt_tokens = ( | |
| fit_messages_to_context( | |
| messages, | |
| MAX_INPUT_TOKENS, | |
| ) | |
| ) | |
| # Never request more than the remaining context. | |
| available_output = ( | |
| CONTEXT_SIZE | |
| - prompt_tokens | |
| - 16 | |
| ) | |
| if available_output < 1: | |
| raise HTTPException( | |
| status_code=400, | |
| detail=( | |
| "Prompt is too large even after " | |
| "context trimming." | |
| ), | |
| ) | |
| max_tokens = min( | |
| max_tokens, | |
| available_output, | |
| ) | |
| print( | |
| f"Request: prompt={prompt_tokens}, " | |
| f"max_tokens={max_tokens}, " | |
| f"stream={stream}" | |
| ) | |
| # ======================================================== | |
| # STREAM | |
| # ======================================================== | |
| if stream: | |
| chat_id = make_id() | |
| created = int(time.time()) | |
| async def stream_generator(): | |
| # ------------------------------------------------ | |
| # Immediate first SSE event. | |
| # ------------------------------------------------ | |
| first_chunk = { | |
| "id": chat_id, | |
| "object": "chat.completion.chunk", | |
| "created": created, | |
| "model": requested_model, | |
| "choices": [ | |
| { | |
| "index": 0, | |
| "delta": { | |
| "role": "assistant", | |
| "content": "", | |
| }, | |
| "finish_reason": None, | |
| } | |
| ], | |
| } | |
| yield ( | |
| "data: " | |
| + json.dumps( | |
| first_chunk, | |
| ensure_ascii=False, | |
| ) | |
| + "\n\n" | |
| ) | |
| try: | |
| async with generation_lock: | |
| result = ( | |
| llm.create_chat_completion( | |
| messages=messages, | |
| temperature=temperature, | |
| top_p=top_p, | |
| max_tokens=max_tokens, | |
| stream=True, | |
| ) | |
| ) | |
| for chunk in result: | |
| choices = chunk.get( | |
| "choices", | |
| [], | |
| ) | |
| if not choices: | |
| continue | |
| choice = choices[0] | |
| delta = choice.get( | |
| "delta", | |
| {}, | |
| ) | |
| content = delta.get( | |
| "content", | |
| "", | |
| ) | |
| if content is None: | |
| content = "" | |
| output_chunk = { | |
| "id": chat_id, | |
| "object": "chat.completion.chunk", | |
| "created": created, | |
| "model": requested_model, | |
| "choices": [ | |
| { | |
| "index": 0, | |
| "delta": { | |
| "content": content, | |
| }, | |
| "finish_reason": choice.get( | |
| "finish_reason" | |
| ), | |
| } | |
| ], | |
| } | |
| yield ( | |
| "data: " | |
| + json.dumps( | |
| output_chunk, | |
| ensure_ascii=False, | |
| ) | |
| + "\n\n" | |
| ) | |
| # ------------------------------------------------ | |
| # Final event. | |
| # ------------------------------------------------ | |
| final_chunk = { | |
| "id": chat_id, | |
| "object": "chat.completion.chunk", | |
| "created": created, | |
| "model": requested_model, | |
| "choices": [ | |
| { | |
| "index": 0, | |
| "delta": {}, | |
| "finish_reason": "stop", | |
| } | |
| ], | |
| } | |
| yield ( | |
| "data: " | |
| + json.dumps( | |
| final_chunk, | |
| ensure_ascii=False, | |
| ) | |
| + "\n\n" | |
| ) | |
| yield "data: [DONE]\n\n" | |
| except Exception as exc: | |
| print( | |
| "STREAM ERROR:", | |
| repr(exc), | |
| ) | |
| error_chunk = { | |
| "id": chat_id, | |
| "object": "chat.completion.chunk", | |
| "created": created, | |
| "model": requested_model, | |
| "choices": [ | |
| { | |
| "index": 0, | |
| "delta": { | |
| "content": ( | |
| "\n[Generation error: " | |
| + str(exc) | |
| + "]" | |
| ) | |
| }, | |
| "finish_reason": "stop", | |
| } | |
| ], | |
| } | |
| yield ( | |
| "data: " | |
| + json.dumps( | |
| error_chunk, | |
| ensure_ascii=False, | |
| ) | |
| + "\n\n" | |
| ) | |
| yield "data: [DONE]\n\n" | |
| return StreamingResponse( | |
| stream_generator(), | |
| media_type="text/event-stream", | |
| headers={ | |
| "Cache-Control": "no-cache, no-transform", | |
| "Connection": "keep-alive", | |
| "X-Accel-Buffering": "no", | |
| }, | |
| ) | |
| # ======================================================== | |
| # NON-STREAM | |
| # ======================================================== | |
| try: | |
| async with generation_lock: | |
| result = ( | |
| llm.create_chat_completion( | |
| messages=messages, | |
| temperature=temperature, | |
| top_p=top_p, | |
| max_tokens=max_tokens, | |
| stream=False, | |
| ) | |
| ) | |
| choices = result.get( | |
| "choices", | |
| [], | |
| ) | |
| if choices: | |
| choice = choices[0] | |
| message = choice.get( | |
| "message", | |
| {}, | |
| ) | |
| output = message.get( | |
| "content", | |
| "", | |
| ) | |
| finish_reason = choice.get( | |
| "finish_reason", | |
| "stop", | |
| ) | |
| else: | |
| output = "" | |
| finish_reason = "stop" | |
| usage = result.get( | |
| "usage", | |
| {}, | |
| ) | |
| return JSONResponse( | |
| content={ | |
| "id": make_id(), | |
| "object": "chat.completion", | |
| "created": int(time.time()), | |
| "model": requested_model, | |
| "choices": [ | |
| { | |
| "index": 0, | |
| "message": { | |
| "role": "assistant", | |
| "content": output, | |
| }, | |
| "finish_reason": finish_reason, | |
| } | |
| ], | |
| "usage": { | |
| "prompt_tokens": usage.get( | |
| "prompt_tokens", | |
| prompt_tokens, | |
| ), | |
| "completion_tokens": usage.get( | |
| "completion_tokens", | |
| 0, | |
| ), | |
| "total_tokens": usage.get( | |
| "total_tokens", | |
| prompt_tokens, | |
| ), | |
| }, | |
| } | |
| ) | |
| except Exception as exc: | |
| print( | |
| "GENERATION ERROR:", | |
| repr(exc), | |
| ) | |
| raise HTTPException( | |
| status_code=500, | |
| detail=f"Generation failed: {str(exc)}", | |
| ) | |
| # ============================================================ | |
| # MAIN | |
| # ============================================================ | |
| if __name__ == "__main__": | |
| import uvicorn | |
| uvicorn.run( | |
| app, | |
| host=HOST, | |
| port=PORT, | |
| ) |