app.py

#!/usr/bin/env python3
"""
DeepCoder Model API Server
Serves the DeepCoder-14B model via FastAPI
"""

import os
import asyncio
import logging
from typing import Optional, Dict, Any
import uvicorn
from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from huggingface_hub import hf_hub_download
import json

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Configuration
MODEL_NAME = os.getenv("MODEL_NAME", "ai/deepcoder-preview")
MODEL_VARIANT = os.getenv("MODEL_VARIANT", "14B-Q4_K_M")
CACHE_DIR = os.getenv("HUGGINGFACE_HUB_CACHE", "/app/cache")
MAX_TOKENS = 131072  # 131K context length
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

app = FastAPI(
    title="DeepCoder API",
    description="AI Code Generation Model API",
    version="1.0.0"
)

# CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Global model variables
tokenizer = None
model = None
model_loaded = False

class CodeRequest(BaseModel):
    prompt: str = Field(..., description="Code generation prompt")
    temperature: float = Field(0.6, ge=0.0, le=2.0, description="Sampling temperature")
    top_p: float = Field(0.95, ge=0.0, le=1.0, description="Top-p sampling")
    max_tokens: int = Field(2048, ge=1, le=8192, description="Maximum tokens to generate")
    stop_sequences: Optional[list] = Field(None, description="Stop sequences")

class CodeResponse(BaseModel):
    generated_code: str
    model_info: Dict[str, Any]
    generation_params: Dict[str, Any]

async def load_model():
    """Load the DeepCoder model and tokenizer"""
    global tokenizer, model, model_loaded
    
    if model_loaded:
        return
    
    try:
        logger.info(f"Loading model: {MODEL_NAME}")
        
        # Load tokenizer
        tokenizer = AutoTokenizer.from_pretrained(
            MODEL_NAME,
            cache_dir=CACHE_DIR,
            trust_remote_code=True
        )
        
        # Load model with appropriate settings for the quantized version
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_NAME,
            cache_dir=CACHE_DIR,
            trust_remote_code=True,
            torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
            device_map="auto" if DEVICE == "cuda" else None,
            load_in_4bit=True if "Q4" in MODEL_VARIANT else False,
        )
        
        if DEVICE == "cpu" and hasattr(model, 'to'):
            model = model.to(DEVICE)
            
        model_loaded = True
        logger.info(f"Model loaded successfully on {DEVICE}")
        
    except Exception as e:
        logger.error(f"Error loading model: {str(e)}")
        raise

@app.on_event("startup")
async def startup_event():
    """Load model on startup"""
    await load_model()

@app.get("/")
async def root():
    return {
        "message": "DeepCoder API",
        "model": MODEL_NAME,
        "variant": MODEL_VARIANT,
        "status": "ready" if model_loaded else "loading"
    }

@app.get("/health")
async def health_check():
    return {
        "status": "healthy" if model_loaded else "loading",
        "model_loaded": model_loaded,
        "device": DEVICE,
        "gpu_available": torch.cuda.is_available()
    }

@app.get("/model/info")
async def model_info():
    """Get model information"""
    if not model_loaded:
        raise HTTPException(status_code=503, detail="Model not loaded yet")
    
    return {
        "model_name": MODEL_NAME,
        "variant": MODEL_VARIANT,
        "max_context_length": MAX_TOKENS,
        "device": DEVICE,
        "model_size": "14B parameters",
        "quantization": "Q4_K_M" if "Q4" in MODEL_VARIANT else "None",
        "benchmarks": {
            "LiveCodeBench_v5_Pass@1": "60.6%",
            "Codeforces_Elo": 1936,
            "Codeforces_Percentile": "95.3",
            "HumanEval+_Accuracy": "92.6%"
        }
    }

@app.post("/generate", response_model=CodeResponse)
async def generate_code(request: CodeRequest):
    """Generate code using the DeepCoder model"""
    if not model_loaded:
        raise HTTPException(status_code=503, detail="Model not loaded yet")
    
    try:
        # Tokenize input
        inputs = tokenizer(
            request.prompt,
            return_tensors="pt",
            truncation=True,
            max_length=MAX_TOKENS - request.max_tokens
        )
        
        if DEVICE == "cuda":
            inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
        
        # Generation parameters
        generation_kwargs = {
            "max_new_tokens": request.max_tokens,
            "temperature": request.temperature,
            "top_p": request.top_p,
            "do_sample": True,
            "pad_token_id": tokenizer.eos_token_id,
        }
        
        if request.stop_sequences:
            generation_kwargs["stop_sequences"] = request.stop_sequences
        
        # Generate
        with torch.no_grad():
            outputs = model.generate(**inputs, **generation_kwargs)
        
        # Decode output
        generated_tokens = outputs[0][inputs["input_ids"].shape[1]:]
        generated_code = tokenizer.decode(generated_tokens, skip_special_tokens=True)
        
        return CodeResponse(
            generated_code=generated_code,
            model_info={
                "model_name": MODEL_NAME,
                "variant": MODEL_VARIANT,
                "device": DEVICE
            },
            generation_params={
                "temperature": request.temperature,
                "top_p": request.top_p,
                "max_tokens": request.max_tokens
            }
        )
        
    except Exception as e:
        logger.error(f"Generation error: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Generation failed: {str(e)}")

@app.post("/chat")
async def chat_completion(request: CodeRequest):
    """Chat-style completion for code assistance"""
    # Add system context for better code generation
    system_prompt = """You are DeepCoder, an expert AI programming assistant. Generate high-quality, well-commented code that follows best practices."""
    
    full_prompt = f"{system_prompt}\n\nUser: {request.prompt}\n\nAssistant:"
    
    # Create modified request with system prompt
    modified_request = CodeRequest(
        prompt=full_prompt,
        temperature=request.temperature,
        top_p=request.top_p,
        max_tokens=request.max_tokens,
        stop_sequences=request.stop_sequences
    )
    
    return await generate_code(modified_request)

if __name__ == "__main__":
    uvicorn.run(
        "app:app",
        host="0.0.0.0",
        port=8000,
        reload=False,
        log_level="info"
    )