Spaces:
Running
Running
File size: 6,757 Bytes
449a0cc 9815afa 449a0cc 9815afa 449a0cc 9815afa 449a0cc 9815afa 449a0cc 9815afa 449a0cc 9815afa 449a0cc 9815afa 449a0cc 9815afa 449a0cc 9815afa 449a0cc 9815afa 449a0cc 9815afa 449a0cc 9815afa 449a0cc 9815afa 449a0cc 9815afa 449a0cc 9815afa 449a0cc |
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 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 |
import os
import torch
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
from typing import Dict, Any
import logging
# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Set cache directories
os.environ['HF_HOME'] = '/app/.cache'
os.environ['TRANSFORMERS_CACHE'] = '/app/.cache/transformers'
os.environ['HF_HUB_CACHE'] = '/app/.cache/hub'
# Inisialisasi API
app = FastAPI(
title="Lyon28 Multi-Model API",
description="API serbaguna untuk 11 model Lyon28"
)
# --- Daftar model dan tugasnya ---
MODEL_MAPPING = {
# Generative Models (Text Generation)
"Tinny-Llama": {"id": "Lyon28/Tinny-Llama", "task": "text-generation"},
"Pythia": {"id": "Lyon28/Pythia", "task": "text-generation"},
"GPT-2": {"id": "Lyon28/GPT-2", "task": "text-generation"},
"GPT-Neo": {"id": "Lyon28/GPT-Neo", "task": "text-generation"},
"Distil_GPT-2": {"id": "Lyon28/Distil_GPT-2", "task": "text-generation"},
"GPT-2-Tinny": {"id": "Lyon28/GPT-2-Tinny", "task": "text-generation"},
# Text-to-Text Model
"T5-Small": {"id": "Lyon28/T5-Small", "task": "text2text-generation"},
# Fill-Mask Models
"Bert-Tinny": {"id": "Lyon28/Bert-Tinny", "task": "fill-mask"},
"Albert-Base-V2": {"id": "Lyon28/Albert-Base-V2", "task": "fill-mask"},
"Distilbert-Base-Uncased": {"id": "Lyon28/Distilbert-Base-Uncased", "task": "fill-mask"},
"Electra-Small": {"id": "Lyon28/Electra-Small", "task": "fill-mask"},
}
# --- Cache untuk menyimpan model yang sudah dimuat ---
PIPELINE_CACHE = {}
def ensure_cache_directory():
"""Pastikan direktori cache ada dan memiliki permission yang benar."""
cache_dirs = [
'/app/.cache',
'/app/.cache/transformers',
'/app/.cache/hub'
]
for cache_dir in cache_dirs:
try:
os.makedirs(cache_dir, exist_ok=True)
os.chmod(cache_dir, 0o755)
logger.info(f"Cache directory {cache_dir} ready")
except Exception as e:
logger.error(f"Failed to create cache directory {cache_dir}: {e}")
def get_pipeline(model_name: str):
"""Fungsi untuk memuat model dari cache atau dari Hub jika belum ada."""
if model_name in PIPELINE_CACHE:
logger.info(f"Mengambil model '{model_name}' dari cache.")
return PIPELINE_CACHE[model_name]
if model_name not in MODEL_MAPPING:
raise HTTPException(status_code=404, detail=f"Model '{model_name}' tidak ditemukan.")
model_info = MODEL_MAPPING[model_name]
model_id = model_info["id"]
task = model_info["task"]
logger.info(f"Memuat model '{model_name}' ({model_id}) untuk tugas '{task}'...")
try:
# Pastikan cache directory siap
ensure_cache_directory()
# Load model dengan error handling yang lebih baik
pipe = pipeline(
task,
model=model_id,
device_map="auto",
cache_dir="/app/.cache/transformers",
trust_remote_code=True # Untuk model custom
)
PIPELINE_CACHE[model_name] = pipe
logger.info(f"Model '{model_name}' berhasil dimuat dan disimpan di cache.")
return pipe
except PermissionError as e:
error_msg = f"Permission error saat memuat model '{model_name}': {str(e)}. Check cache directory permissions."
logger.error(error_msg)
raise HTTPException(status_code=500, detail=error_msg)
except Exception as e:
error_msg = f"Gagal memuat model '{model_name}': {str(e)}. Common causes: 1) another user is downloading the same model (please wait); 2) a previous download was canceled and the lock file needs manual removal."
logger.error(error_msg)
raise HTTPException(status_code=500, detail=error_msg)
# --- Definisikan struktur request dari user ---
class InferenceRequest(BaseModel):
model_name: str # Nama kunci dari MODEL_MAPPING, misal: "Tinny-Llama"
prompt: str
parameters: Dict[str, Any] = {} # Parameter tambahan seperti max_length, temperature, dll.
@app.get("/")
def read_root():
"""Endpoint untuk mengecek status API dan daftar model yang tersedia."""
return {
"status": "API is running!",
"available_models": list(MODEL_MAPPING.keys()),
"cached_models": list(PIPELINE_CACHE.keys()),
"cache_info": {
"HF_HOME": os.environ.get('HF_HOME'),
"TRANSFORMERS_CACHE": os.environ.get('TRANSFORMERS_CACHE'),
"HF_HUB_CACHE": os.environ.get('HF_HUB_CACHE')
}
}
@app.get("/health")
def health_check():
"""Health check endpoint."""
return {"status": "healthy", "cached_models": len(PIPELINE_CACHE)}
@app.post("/invoke")
def invoke_model(request: InferenceRequest):
"""Endpoint utama untuk melakukan inferensi pada model yang dipilih."""
try:
# Ambil atau muat pipeline model
pipe = get_pipeline(request.model_name)
# Gabungkan prompt dengan parameter tambahan
result = pipe(request.prompt, **request.parameters)
return {
"model_used": request.model_name,
"prompt": request.prompt,
"parameters": request.parameters,
"result": result
}
except HTTPException as e:
# Meneruskan error yang sudah kita definisikan
raise e
except Exception as e:
# Menangkap error lain yang mungkin terjadi saat inferensi
logger.error(f"Inference error: {str(e)}")
raise HTTPException(status_code=500, detail=f"Terjadi error saat inferensi: {str(e)}")
@app.delete("/cache/{model_name}")
def clear_model_cache(model_name: str):
"""Endpoint untuk menghapus model dari cache."""
if model_name in PIPELINE_CACHE:
del PIPELINE_CACHE[model_name]
logger.info(f"Model '{model_name}' removed from cache")
return {"status": "success", "message": f"Model '{model_name}' removed from cache"}
else:
raise HTTPException(status_code=404, detail=f"Model '{model_name}' tidak ada di cache")
# Startup event dengan error handling yang lebih baik
@app.on_event("startup")
async def startup_event():
logger.info("API startup: Melakukan warm-up dengan memuat satu model awal...")
# Pastikan cache directory siap
ensure_cache_directory()
try:
# Coba model yang paling kecil terlebih dahulu
get_pipeline("GPT-2-Tinny")
logger.info("Warm-up berhasil!")
except Exception as e:
logger.warning(f"Gagal melakukan warm-up: {e}")
logger.info("API tetap berjalan, model akan dimuat saat diperlukan.") |