Rename main.py to app.py

main.py → app.py

@@ -4,6 +4,7 @@ from transformers import pipeline
 import torch
 from fastapi.middleware.cors import CORSMiddleware
 from typing import Dict, Any
+import os # Import os module
 
 # Inisialisasi aplikasi FastAPI
 app = FastAPI(
@@ -53,13 +54,18 @@ def get_task(model_id: str) -> str:
     for task, models in TASK_MAP.items():
         if model_id in models:
             return task
-    return "text-generation"
+    # Default to text-generation if not found (or raise an error)
+    return "text-generation" 
 
 # Event startup untuk inisialisasi model
 @app.on_event("startup")
 async def load_models():
     app.state.pipelines = {}
     print("🟢 Semua model siap digunakan!")
+    # Menyetel HF_HOME untuk mengatasi masalah izin cache
+    os.environ['HF_HOME'] = '/tmp/.cache/huggingface'
+    os.makedirs(os.environ['HF_HOME'], exist_ok=True)
+
 
 # Endpoint utama
 @app.get("/")
@@ -96,11 +102,14 @@ async def health_check():
 @app.post("/inference/{model_id}")
 async def model_inference(model_id: str, request: InferenceRequest):
     try:
+        # Pastikan model_id dalam lowercase agar sesuai dengan MODEL_MAP
+        model_id = model_id.lower() 
+
         # Validasi model ID
         if model_id not in MODEL_MAP:
             raise HTTPException(
                 status_code=404,
-                detail=f"Model {model_id} tidak ditemukan. Cek /models untuk list model yang tersedia."
+                detail=f"Model '{model_id}' tidak ditemukan. Cek /models untuk list model yang tersedia."
             )
 
         # Dapatkan task yang sesuai
@@ -108,11 +117,18 @@ async def model_inference(model_id: str, request: InferenceRequest):
         
         # Load model jika belum ada di memory
         if model_id not in app.state.pipelines:
+            print(f"⏳ Memuat model {model_id} untuk task {task}...")
+            # Menggunakan device=-1 (CPU) sebagai default yang aman
+            # Jika Anda yakin Hugging Face Space Anda memiliki GPU, gunakan device=0
+            device_to_use = 0 if torch.cuda.is_available() else -1
+            # Menyesuaikan dtype berdasarkan device
+            dtype_to_use = torch.float16 if torch.cuda.is_available() else torch.float32
+
             app.state.pipelines[model_id] = pipeline(
                 task=task,
                 model=MODEL_MAP[model_id],
-                device=0 if torch.cuda.is_available() else -1,
-                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+                device=device_to_use,
+                torch_dtype=dtype_to_use
             )
             print(f"✅ Model {model_id} berhasil dimuat!")
 
@@ -128,6 +144,7 @@ async def model_inference(model_id: str, request: InferenceRequest):
             )[0]['generated_text']
         
         elif task == "text-classification":
+            # Untuk text-classification, output adalah list of dict, kita ambil yang pertama
             output = pipe(request.text)[0]
             result = {
                 "label": output['label'],
@@ -135,19 +152,33 @@ async def model_inference(model_id: str, request: InferenceRequest):
             }
         
         elif task == "text2text-generation":
+            # Untuk text2text-generation, output juga list of dict
             result = pipe(
                 request.text,
                 max_length=request.max_length
             )[0]['generated_text']
+        
+        else:
+            # Fallback untuk task yang tidak terduga, meski harusnya terhandle oleh get_task
+            raise HTTPException(
+                status_code=500,
+                detail=f"Tugas ({task}) untuk model {model_id} tidak didukung atau tidak dikenali."
+            )
 
         return {"result": result}
     
     except Exception as e:
+        # Log error lebih detail untuk debugging
+        print(f"‼️ Error saat memproses model {model_id}: {e}")
+        import traceback
+        traceback.print_exc() # Mencetak full traceback ke log
+
         raise HTTPException(
             status_code=500,
-            detail=f"Error processing request: {str(e)}"
+            detail=f"Error processing request: {str(e)}. Cek log server untuk detail."
         )
 
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+# Ini tidak perlu dijalankan secara langsung di Hugging Face Spaces karena Uvicorn akan menjalankannya
+# if __name__ == "__main__":
+#     import uvicorn
+#     uvicorn.run(app, host="0.0.0.0", port=7860)