loleg/fastapi-apertus

1from contextlib import asynccontextmanager2from fastapi import FastAPI, HTTPException3from fastapi.middleware.cors import CORSMiddleware4from pydantic import BaseModel, ValidationError5from typing import List, Optional6 7from torch import cuda8from transformers import (9    AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig10)11 12from hashlib import sha25613from huggingface_hub import login14from dotenv import load_dotenv15from datetime import datetime16 17import os18import uvicorn19import time20import logging21 22# Configure logging23logging.basicConfig(level=logging.INFO)24logger = logging.getLogger(__name__)25 26# Required for access to a gated model27load_dotenv()28hf_token = os.getenv("HF_TOKEN", None)29if hf_token is not None:30    login(token=hf_token)31 32# Configurable model identifier33model_name = os.getenv("HF_MODEL", "swiss-ai/Apertus-8B-Instruct-2509")34model_quantization = int(os.getenv("QUANTIZE", 0)) # 8, 4, 0=default35 36# Configure max tokens37MAX_NEW_TOKENS = 409638 39# Load base prompt from a text file40system_prompt = None41if int(os.getenv("USE_SYSTEM_PROMPT", 1)):42    with open('system_prompt.md', 'r') as file:43        system_prompt = file.read()44 45# Keep data in session46model = None47tokenizer = None48 49class TextInput(BaseModel):50    text: str = ""51    min_length: int = 352    # Apertus by default supports a context length up to 65,536 tokens.53    max_length: int = 6553654 55class ModelResponse(BaseModel):56    text: str57    confidence: float58    processing_time: float59 60class ChatMessage(BaseModel):61    role: str = "user"62    content: str = ""63 64class Completion(BaseModel):65    model: str = "apertus"66    messages: List[ChatMessage]67    max_tokens: Optional[int] = 51268    temperature: Optional[float] = 0.169    top_p: Optional[float] = 0.970 71@asynccontextmanager72async def lifespan(app: FastAPI):73    """Load the transformer model on startup"""74    global model, tokenizer75    try:76        logger.info(f"Loading model: {model_name}")77 78        # Automatically select device based on availability79        device = "cuda" if cuda.is_available() else "cpu"80 81        # load the tokenizer and the model82        tokenizer = AutoTokenizer.from_pretrained(model_name)83 84        # Use a quantization setting85        bnb_config = None86        if model_quantization == 8:87            bnb_config = BitsAndBytesConfig(load_in_8bit=True)88        elif model_quantization == 4:89            bnb_config = BitsAndBytesConfig(load_in_4bit=True)90        if bnb_config is not None:91            model = AutoModelForCausalLM.from_pretrained(92                model_name,93                device_map="auto",              # Automatically splits model across CPU/GPU94                offload_folder="offload",       # Temporary offload to disk95                low_cpu_mem_usage=True,         # Avoids unnecessary CPU memory duplication96                quantization_config=bnb_config, # To reduce memory and overhead97            )98        else:99            model = AutoModelForCausalLM.from_pretrained(100                model_name,101                device_map="auto",              # Automatically splits model across CPU/GPU102                offload_folder="offload",       # Temporary offload to disk103            )104        logger.info(f"Model loaded successfully! ({device})")105    except Exception as e:106        logger.error(f"Failed to load model: {e}")107        raise e108    # Release resources when the app is stopped109    yield110    del model111    del tokenizer112    cuda.empty_cache()113 114 115# Setup our app116app = FastAPI(117    title="Apertus API",118    description="REST API for serving Apertus models via Hugging Face transformers",119    version="0.1.0",120    docs_url="/",121    lifespan=lifespan122)123 124app.add_middleware(125    CORSMiddleware,126    allow_origins=["*"],127    allow_credentials=True,128    allow_methods=["*"],129    allow_headers=["*"],130)131 132 133def fit_to_length(text, min_length=3, max_length=100):134    """Truncate text if too long."""135    text = text[:max_length]136    if len(text) == max_length:137        logger.warning("Warning: text truncated")138    if len(text) < min_length:139        logger.warning("Warning: empty text, aborting")140        return None141    return text142 143def get_completion_text(messages_think: List[ChatMessage]):144    txt = ""145    for cm in messages_think:146        txt = " ".join((txt, cm.content))147    return txt148 149 150def get_message_id(txt: str):151    return sha256(str(txt).encode()).hexdigest()152 153 154def get_model_reponse(messages_think: List[ChatMessage]):155    """Process the text content."""156 157    # Apply the system template158    has_system = False159    for m in messages_think:160        if m.role == 'system':161            has_system = True162    if not has_system and system_prompt:163        cm = ChatMessage(role='system', content=system_prompt)164        messages_think.insert(0, cm)165    print(messages_think)166 167    # Prepare the model input168    text = tokenizer.apply_chat_template(169        messages_think,170        tokenize=False,171        add_generation_prompt=True,172        top_p=0.9,173        temperature=0.8,174    )175    model_inputs = tokenizer(176        [text], 177        return_tensors="pt",178        add_special_tokens=False179    ).to(model.device)180 181    # Generate the output182    generated_ids = model.generate(183        **model_inputs,184        max_new_tokens=MAX_NEW_TOKENS185    )186 187    # Get and decode the output188    output_ids = generated_ids[0][len(model_inputs.input_ids[0]) :]189    190    # Decode the text message191    return tokenizer.decode(output_ids, skip_special_tokens=True)192 193 194@app.post("/v1/models/apertus")195async def completion(data: Completion):196    """Generate an OpenAPI-style completion"""197    if model is None or tokenizer is None:198        raise HTTPException(status_code=503, detail="Model not loaded")199 200    try:201        mt = data.messages202        text = get_completion_text(mt)203        result = get_model_reponse(mt)204 205        # Standard formatted object206        return {207            "id": get_message_id(text),208            "object": "chat.completion",209            "created": time.time(),210            "model": data.model,211            "choices": [{212                "message": ChatMessage(role="assistant", content=result)213            }],214            "usage": {215                "prompt_tokens": len(text),216                "completion_tokens": len(result),217                "total_tokens": len(text) + len(result)218            }219        }220    except Exception as e:221        logger.warning(e)222        raise HTTPException(status_code=400, detail="Could not process") from e223 224 225@app.get("/predict", response_model=ModelResponse)226async def predict(q: str):227    """Generate a model response for input text"""228    if model is None or tokenizer is None:229        raise HTTPException(status_code=503, detail="Model not loaded")230 231    try:232        start_time = time.time()233 234        input_data = TextInput(text=q)235 236        text = fit_to_length(input_data.text, input_data.min_length, input_data.max_length)237 238        messages_think = [239            {"role": "user", "content": text}240        ]241        result = get_model_reponse(messages_think)242 243        # Checkpoint244        processing_time = time.time() - start_time245 246        return ModelResponse(247            text=result, #['label'],248            confidence=0, #result['score'],249            processing_time=processing_time250        )251 252    except Exception as e:253        logger.warning(e)254        raise HTTPException(status_code=500, detail="Evaluation failed")255 256@app.get("/health")257async def health_check():258    """Health check and basic configuration"""259    return {260        "status": "healthy",261        "model_loaded": model is not None,262        "gpu_available": cuda.is_available()263    }264 265if __name__=='__main__':266    uvicorn.run('app:app', reload=True)267