Spaces:
Runtime error
Runtime error
File size: 4,352 Bytes
5eb40a3 4f97961 5eb40a3 1303c4e 98cad23 1303c4e 5eb40a3 1303c4e 98cad23 64d30b7 98cad23 5eb40a3 1303c4e 5eb40a3 64d30b7 5eb40a3 4f97961 5eb40a3 1303c4e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 |
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from torch import cuda
from transformers import AutoModelForCausalLM, AutoTokenizer
from huggingface_hub import login
from dotenv import load_dotenv
import os
import logging
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Required for access to a gated model
load_dotenv()
hf_token = os.getenv("HF_TOKEN", None)
if hf_token is not None:
login(token=hf_token)
# Keep data in session
model = None
tokenizer = None
class TextInput(BaseModel):
text: str
min_length: int = 3
# Apertus by default supports a context length up to 65,536 tokens.
max_length: int = 65536
class ModelResponse(BaseModel):
text: str
confidence: float
processing_time: float
@asynccontextmanager
async def lifespan(app: FastAPI):
"""Load the transformer model on startup"""
global model, tokenizer
try:
logger.info("Loading sentiment analysis model...")
# TODO: make this configurable
model_name = "swiss-ai/Apertus-8B-Instruct-2509"
# Automatically select device based on availability
device = "cuda" if cuda.is_available() else "cpu"
# load the tokenizer and the model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
model_name,
device_map="auto", # Automatically splits model across CPU/GPU
low_cpu_mem_usage=True, # Avoids unnecessary CPU memory duplication
offload_folder="offload", # Temporary offload to disk
)
#.to(device)
logger.info(f"Model loaded successfully! ({device})")
except Exception as e:
logger.error(f"Failed to load model: {e}")
raise e
# Release resources when the app is stopped
yield
model.clear()
tokenizer.clear()
# Setup our app
app = FastAPI(
title="Apertus API",
description="REST API for serving Apertus models via Hugging Face transformers",
version="0.1.0",
lifespan=lifespan
)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
@app.get("/predict", response_model=ModelResponse)
async def predict(q: str):
"""Generate a model response for input text"""
if model is None or tokenizer is None:
raise HTTPException(status_code=503, detail="Model not loaded")
try:
import time
start_time = time.time()
input_data = TextInput(text=q)
# Truncate text if too long
text = input_data.text[:input_data.max_length]
if len(text) == input_data.max_length:
logger.warning("Warning: text truncated")
if len(text) < input_data.min_length:
logger.warning("Warning: empty text, aborting")
return None
# Prepare the model input
messages_think = [
{"role": "user", "content": text}
]
text = tokenizer.apply_chat_template(
messages_think,
tokenize=False,
add_generation_prompt=True,
)
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
# Generate the output
generated_ids = model.generate(**model_inputs, max_new_tokens=32768)
# Get and decode the output
output_ids = generated_ids[0][len(model_inputs.input_ids[0]) :]
result = tokenizer.decode(output_ids, skip_special_tokens=True)
# Checkpoint
processing_time = time.time() - start_time
return ModelResponse(
text=result['label'],
confidence=result['score'],
processing_time=processing_time
)
except Exception as e:
logger.error(f"Evaluation error: {e}")
raise HTTPException(status_code=500, detail="Evaluation failed")
@app.get("/health")
async def health_check():
"""Health check and basic configuration"""
return {
"status": "healthy",
"model_loaded": model is not None,
"gpu_available": cuda.is_available()
}
@app.get("/")
def read_root():
return PlainTextResponse('Habemus Apertus')
|