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.")