Create app.py

app.py

@@ -0,0 +1,300 @@
+import os
+import json
+import logging
+from pathlib import Path
+from typing import Optional
+from datetime import datetime
+
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import StreamingResponse, JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+import uvicorn
+
+try:
+    from llama_cpp import Llama
+except ImportError:
+    print("Installing llama-cpp-python...")
+    os.system("pip install llama-cpp-python -q")
+    from llama_cpp import Llama
+
+# ============================================================================
+# SETUP & CONFIG
+# ============================================================================
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+app = FastAPI(title="Qwen 3B Chat API", version="1.0.0")
+
+# Enable CORS for all origins (or restrict as needed)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Persistent storage path (HF Space mounts this)
+STORAGE_DIR = Path("/data")  # Default HF persistent storage mount point
+STORAGE_DIR.mkdir(exist_ok=True)
+
+MODEL_CACHE_DIR = STORAGE_DIR / "models"
+MODEL_CACHE_DIR.mkdir(exist_ok=True)
+
+# Model configuration
+MODEL_REPO = "Qwen/Qwen2.5-3B-Instruct-GGUF"
+MODEL_FILE = "qwen2.5-3b-instruct-q4_k_m.gguf"
+MODEL_PATH = MODEL_CACHE_DIR / MODEL_FILE
+
+# Global model instance
+llm = None
+
+# ============================================================================
+# REQUEST/RESPONSE MODELS
+# ============================================================================
+
+class ChatMessage(BaseModel):
+    role: str  # "user" or "assistant"
+    content: str
+
+class ChatRequest(BaseModel):
+    messages: list[ChatMessage]
+    max_tokens: int = 512
+    temperature: float = 0.7
+    top_p: float = 0.9
+    stream: bool = False
+
+class ChatResponse(BaseModel):
+    id: str
+    created: int
+    model: str
+    choices: list[dict]
+    usage: dict
+
+# ============================================================================
+# MODEL LOADING & INITIALIZATION
+# ============================================================================
+
+def download_model():
+    """Download model from HF Hub to persistent storage if not already cached."""
+    if MODEL_PATH.exists():
+        logger.info(f"Model found at {MODEL_PATH}")
+        return
+    
+    logger.info(f"Downloading {MODEL_REPO}/{MODEL_FILE} to {MODEL_PATH}...")
+    
+    try:
+        from huggingface_hub import hf_hub_download
+        
+        hf_hub_download(
+            repo_id=MODEL_REPO,
+            filename=MODEL_FILE,
+            local_dir=MODEL_CACHE_DIR,
+            local_dir_use_symlinks=False
+        )
+        logger.info(f"Model downloaded successfully to {MODEL_PATH}")
+    except Exception as e:
+        logger.error(f"Failed to download model: {e}")
+        raise
+
+def load_model():
+    """Load the GGUF model with llama-cpp-python."""
+    global llm
+    
+    if llm is not None:
+        logger.info("Model already loaded")
+        return
+    
+    if not MODEL_PATH.exists():
+        download_model()
+    
+    logger.info(f"Loading model from {MODEL_PATH}...")
+    
+    try:
+        llm = Llama(
+            model_path=str(MODEL_PATH),
+            n_gpu_layers=-1,  # Use all GPU layers if GPU available, otherwise CPU
+            n_ctx=2048,       # Context window
+            n_threads=4,      # CPU threads (match Space CPU count)
+            verbose=False,
+        )
+        logger.info("Model loaded successfully")
+    except Exception as e:
+        logger.error(f"Failed to load model: {e}")
+        raise
+
+@app.on_event("startup")
+async def startup_event():
+    """Load model on startup."""
+    load_model()
+
+# ============================================================================
+# UTILITY FUNCTIONS
+# ============================================================================
+
+def format_messages_for_prompt(messages: list[ChatMessage]) -> str:
+    """Format chat messages into a prompt for Qwen."""
+    # Qwen2.5 uses a simple chat format
+    prompt = ""
+    for msg in messages:
+        if msg.role == "user":
+            prompt += f"User: {msg.content}\n"
+        elif msg.role == "assistant":
+            prompt += f"Assistant: {msg.content}\n"
+    prompt += "Assistant:"
+    return prompt
+
+def parse_stream(output: str) -> str:
+    """Extract just the response content from model output."""
+    # Model may include "User: " prefix if continuing conversation
+    # Clean up for cleaner response
+    if "User:" in output:
+        output = output.split("User:")[0]
+    return output.strip()
+
+# ============================================================================
+# API ENDPOINTS
+# ============================================================================
+
+@app.get("/health")
+async def health():
+    """Health check endpoint."""
+    return {"status": "healthy", "model_loaded": llm is not None}
+
+@app.post("/v1/chat/completions")
+async def chat_completions(request: ChatRequest):
+    """
+    OpenAI-compatible chat completion endpoint.
+    Supports both streaming and non-streaming responses.
+    """
+    if llm is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+    
+    try:
+        prompt = format_messages_for_prompt(request.messages)
+        
+        if request.stream:
+            return StreamingResponse(
+                stream_chat(prompt, request),
+                media_type="text/event-stream"
+            )
+        else:
+            # Non-streaming response
+            output = llm(
+                prompt,
+                max_tokens=request.max_tokens,
+                temperature=request.temperature,
+                top_p=request.top_p,
+                stop=["User:", "\n\n"],
+            )
+            
+            response_text = output["choices"][0]["text"].strip()
+            
+            return ChatResponse(
+                id=f"chatcmpl-{datetime.now().timestamp()}",
+                created=int(datetime.now().timestamp()),
+                model="qwen2.5-3b-instruct",
+                choices=[
+                    {
+                        "index": 0,
+                        "message": {
+                            "role": "assistant",
+                            "content": response_text
+                        },
+                        "finish_reason": "stop"
+                    }
+                ],
+                usage={
+                    "prompt_tokens": len(prompt.split()),
+                    "completion_tokens": len(response_text.split()),
+                    "total_tokens": len(prompt.split()) + len(response_text.split())
+                }
+            )
+    
+    except Exception as e:
+        logger.error(f"Error in chat completions: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+async def stream_chat(prompt: str, request: ChatRequest):
+    """Stream tokens from the model."""
+    try:
+        # Use llama-cpp-python's streaming capability
+        for chunk in llm(
+            prompt,
+            max_tokens=request.max_tokens,
+            temperature=request.temperature,
+            top_p=request.top_p,
+            stop=["User:", "\n\n"],
+            stream=True,
+        ):
+            token = chunk["choices"][0]["text"]
+            if token:
+                # OpenAI-compatible streaming format
+                data = {
+                    "id": f"chatcmpl-{datetime.now().timestamp()}",
+                    "created": int(datetime.now().timestamp()),
+                    "model": "qwen2.5-3b-instruct",
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {"role": "assistant", "content": token},
+                            "finish_reason": None
+                        }
+                    ]
+                }
+                yield f"data: {json.dumps(data)}\n\n"
+        
+        # Send final chunk
+        yield f"data: [DONE]\n\n"
+    
+    except Exception as e:
+        logger.error(f"Error in stream_chat: {e}")
+        yield f"data: {json.dumps({'error': str(e)})}\n\n"
+
+@app.post("/chat")
+async def chat_simple(request: ChatRequest):
+    """Simple chat endpoint (returns full response at once)."""
+    response = await chat_completions(request)
+    
+    # If streaming response, convert to regular response
+    if isinstance(response, StreamingResponse):
+        raise HTTPException(status_code=400, detail="Disable streaming for this endpoint")
+    
+    return response
+
+@app.get("/")
+async def root():
+    """Root endpoint with API documentation."""
+    return {
+        "name": "Qwen 3B Chat API",
+        "version": "1.0.0",
+        "model": "Qwen2.5-3B-Instruct",
+        "endpoints": {
+            "/health": "GET - Health check",
+            "/v1/chat/completions": "POST - OpenAI-compatible chat completions",
+            "/chat": "POST - Simple chat endpoint",
+            "/docs": "Interactive API documentation (Swagger UI)",
+        },
+        "example_request": {
+            "messages": [
+                {"role": "user", "content": "Hello, how are you?"}
+            ],
+            "max_tokens": 512,
+            "temperature": 0.7,
+            "stream": False
+        }
+    }
+
+# ============================================================================
+# RUN
+# ============================================================================
+
+if __name__ == "__main__":
+    uvicorn.run(
+        app,
+        host="0.0.0.0",
+        port=7860,  # HF Spaces default port
+        log_level="info"
+    )