Update app.py

app.py

@@ -1,20 +1,21 @@
-# app.py untuk Hugging Face Spaces
-from flask import Flask, jsonify, request, render_template_string
+# app.py
+from flask import Flask, jsonify, request
 from flask_cors import CORS
 from transformers import pipeline
 import logging
-import gc
 import torch
+import os # Untuk mendapatkan environment variables, misalnya di Hugging Face Spaces
 
 app = Flask(__name__)
-CORS(app)
+CORS(app) # Mengaktifkan CORS untuk mengizinkan permintaan dari frontend Anda
 
-# Setup logging
-logging.basicConfig(level=logging.INFO)
+# --- Setup Logging ---
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-# Global models storage
-models = {}
+# --- Konfigurasi Model dan Informasi ---
+# Tambahkan 'hf_model_name' jika nama model di Hugging Face berbeda dari ID yang Anda inginkan.
+# Jika nama model di Hugging Face sama, tidak perlu 'hf_model_name'.
 model_info = {
     "Albert-Base-V2": {"task": "fill-mask", "description": "BERT-based model for masked language modeling"},
     "GPT-2": {"task": "text-generation", "description": "GPT-2 model for text generation"},
@@ -24,117 +25,174 @@ model_info = {
     "Bert-Tinny": {"task": "fill-mask", "description": "Tiny BERT model"},
     "Distilbert-Base-Uncased": {"task": "fill-mask", "description": "Distilled BERT model"},
     "Pythia": {"task": "text-generation", "description": "Pythia language model"},
-    "T5-Small": {"task": "text2text-generation", "description": "Small T5 model"},
+    "T5-Small": {"task": "text2text-generation", "description": "Small T5 model", "hf_model_name": "t5-small"},
     "GPT-Neo": {"task": "text-generation", "description": "GPT-Neo model"},
     "Distil-GPT-2": {"task": "text-generation", "description": "Distilled GPT-2 model"}
 }
 
-def load_models():
-    """Load all models at startup"""
-    logger.info("Loading all models...")
-    
-    for model_name, info in model_info.items():
+# --- Penyimpanan Model Global (untuk Lazy Loading) ---
+models = {}
+
+# --- Fungsi Utility untuk Model Lazy Loading ---
+def get_model_pipeline(model_name):
+    """
+    Memuat model hanya jika belum dimuat (lazy loading).
+    Mengembalikan pipeline model yang diminta.
+    """
+    if model_name not in models:
+        logger.info(f"Model '{model_name}' belum dimuat. Memuat sekarang...")
+        if model_name not in model_info:
+            logger.error(f"Informasi model '{model_name}' tidak ditemukan di model_info.")
+            raise ValueError(f"Model '{model_name}' tidak dikenal.")
+
+        info = model_info[model_name]
         try:
-            logger.info(f"Loading {model_name}...")
-            model_path = f"Lyon28/{model_name}"
-            
-            # Load model with appropriate task
+            # Gunakan 'hf_model_name' jika disediakan, jika tidak, gunakan model_name dengan prefix 'Lyon28/'
+            hf_model_path = info.get("hf_model_name", f"Lyon28/{model_name}")
+
+            # Explicitly set device to "cpu" for CPU-only environments
             models[model_name] = pipeline(
-                info["task"], 
-                model=model_path,
-                device=-1,  # CPU only
-                torch_dtype=torch.float32
+                info["task"],
+                model=hf_model_path,
+                device="cpu", # Penting: Pastikan ini "cpu" jika Anda tidak punya GPU
+                torch_dtype=torch.float32 # Tetap float32 untuk performa terbaik di CPU
             )
-            logger.info(f"✅ {model_name} loaded successfully")
-            
+            logger.info(f"✅ Model '{model_name}' (Path: {hf_model_path}) berhasil dimuat.")
         except Exception as e:
-            logger.error(f"❌ Failed to load {model_name}: {str(e)}")
-            
-    logger.info(f"Loaded {len(models)}/{len(model_info)} models")
+            logger.error(f"❌ Gagal memuat model '{model_name}' (Path: {hf_model_path}): {str(e)}", exc_info=True)
+            raise RuntimeError(f"Gagal memuat model: {model_name}. Detail: {str(e)}") from e
+    return models[model_name]
 
-# Load models on startup
-load_models()
+# --- Rute API ---
 
 @app.route('/')
 def home():
-    """Simple API status - no HTML interface"""
+    """Endpoint root untuk status API."""
     return jsonify({
-        "message": "Lyon28's AI Models API",
+        "message": "Flask API untuk Model Hugging Face",
         "status": "online",
-        "total_models": len(model_info),
-        "loaded_models": len(models),
-        "endpoints": {
-            "models_list": "/api/models",
-            "health_check": "/health",
-            "prediction": "/api/{model_name}"
-        },
-        "available_models": list(models.keys())
+        "loaded_models_count": len(models),
+        "available_model_configs": list(model_info.keys()),
+        "info": "Gunakan /api/models untuk daftar model yang tersedia."
     })
 
 @app.route('/api/models', methods=['GET'])
-def list_models():
-    """Get list of available models"""
-    available_models = []
-    for name, info in model_info.items():
-        available_models.append({
-            "name": name,
+def list_available_models():
+    """Mengembalikan daftar semua model yang dikonfigurasi, termasuk status muatan."""
+    available_models_data = [
+        {
+            "id": name,
+            "name": info["description"],
             "task": info["task"],
-            "description": info["description"],
-            "status": "ready" if name in models else "failed",
+            "status": "loaded" if name in models else "not_loaded", # Menunjukkan apakah sudah dimuat via lazy loading
             "endpoint": f"/api/{name}"
-        })
-    
+        }
+        for name, info in model_info.items()
+    ]
     return jsonify({
-        "total": len(model_info),
-        "loaded": len(models),
-        "models": available_models
+        "total_configured_models": len(model_info),
+        "currently_loaded_models": len(models),
+        "models": available_models_data
     })
 
-@app.route('/api/<model_name>', methods=['POST'])
-def predict(model_name):
-    """Main prediction endpoint"""
-    if model_name not in models:
-        return jsonify({
-            "error": f"Model '{model_name}' not available. Available models: {list(models.keys())}"
-        }), 404
-    
+@app.route('/api/<model_id>', methods=['POST'])
+def predict_with_model(model_id):
+    """
+    Endpoint utama untuk prediksi model.
+    Menerima 'inputs' (teks) dan 'parameters' (dictionary) opsional.
+    """
+    logger.info(f"Menerima permintaan untuk model: {model_id}")
+    if model_id not in model_info:
+        logger.warning(f"Permintaan untuk model tidak dikenal: {model_id}")
+        return jsonify({"error": f"Model '{model_id}' tidak dikenal. Lihat /api/models untuk daftar yang tersedia."}), 404
+
     try:
+        model_pipeline = get_model_pipeline(model_id) # Memuat model jika belum ada
+        model_task = model_info[model_id]["task"]
+
         data = request.json
         inputs = data.get('inputs', '')
-        parameters = data.get('parameters', {})
-        
+        parameters = data.get('parameters', {}) # Default ke dictionary kosong jika tidak ada
+
         if not inputs:
-            return jsonify({"error": "No inputs provided"}), 400
-        
-        # Get model
-        model = models[model_name]
-        
-        # Generate prediction
-        if parameters:
-            result = model(inputs, **parameters)
+            return jsonify({"error": "Input 'inputs' tidak boleh kosong."}), 400
+
+        logger.info(f"Inferensi: Model='{model_id}', Task='{model_task}', Input='{inputs[:100]}...', Params='{parameters}'")
+
+        result = []
+        # --- Penanganan Parameter dan Inferensi berdasarkan Tipe Tugas ---
+        if model_task == "text-generation":
+            # Default parameters for text-generation
+            gen_params = {
+                "max_new_tokens": parameters.get("max_new_tokens", 150), # Lebih banyak token untuk roleplay
+                "temperature": parameters.get("temperature", 0.7),
+                "do_sample": parameters.get("do_sample", True),
+                "return_full_text": parameters.get("return_full_text", False), # Sangat penting untuk chatbot
+                "num_return_sequences": parameters.get("num_return_sequences", 1),
+                "top_k": parameters.get("top_k", 50),
+                "top_p": parameters.get("top_p", 0.95),
+                "repetition_penalty": parameters.get("repetition_penalty", 1.2), # Mencegah pengulangan
+            }
+            result = model_pipeline(inputs, **gen_params)
+
+        elif model_task == "fill-mask":
+            mask_params = {
+                "top_k": parameters.get("top_k", 5)
+            }
+            result = model_pipeline(inputs, **mask_params)
+
+        elif model_task == "text2text-generation": # Misalnya untuk T5
+            t2t_params = {
+                "max_new_tokens": parameters.get("max_new_tokens", 150),
+                "temperature": parameters.get("temperature", 0.7),
+                "do_sample": parameters.get("do_sample", True),
+            }
+            result = model_pipeline(inputs, **t2t_params)
+
+        else:
+            # Fallback for other tasks or if no specific parameters are needed
+            result = model_pipeline(inputs, **parameters)
+
+        # --- Konsistensi Format Output ---
+        response_output = {}
+        if model_task == "text-generation" or model_task == "text2text-generation":
+            if result and len(result) > 0 and 'generated_text' in result[0]:
+                response_output['text'] = result[0]['generated_text'].strip()
+            else:
+                response_output['text'] = "[Tidak ada teks yang dihasilkan atau format tidak sesuai.]"
+        elif model_task == "fill-mask":
+            response_output['predictions'] = [
+                {"sequence": p.get('sequence', ''), "score": p.get('score', 0.0), "token_str": p.get('token_str', '')}
+                for p in result
+            ]
         else:
-            result = model(inputs)
-        
-        return jsonify({
-            "model": model_name,
-            "inputs": inputs,
-            "outputs": result,
-            "parameters": parameters
-        })
-        
+            # Untuk jenis tugas lain, kembalikan hasil mentah
+            response_output = result
+
+        logger.info(f"Inferensi berhasil untuk '{model_id}'. Output singkat: '{str(response_output)[:200]}'")
+        return jsonify({"model": model_id, "inputs": inputs, "outputs": response_output})
+
+    except ValueError as ve:
+        # Error yang berasal dari get_model_pipeline atau validasi input
+        logger.error(f"Validasi atau konfigurasi error untuk model '{model_id}': {str(ve)}")
+        return jsonify({"error": str(ve), "message": "Kesalahan konfigurasi atau input model."}), 400
+    except RuntimeError as re:
+        # Error saat memuat model
+        logger.error(f"Error runtime saat memuat model '{model_id}': {str(re)}")
+        return jsonify({"error": str(re), "message": "Model gagal dimuat."}), 503 # Service Unavailable
     except Exception as e:
-        logger.error(f"Error with model {model_name}: {str(e)}")
-        return jsonify({"error": str(e)}), 500
+        # Catch all other unexpected errors during prediction
+        logger.error(f"Terjadi kesalahan tak terduga saat memprediksi dengan model '{model_id}': {str(e)}", exc_info=True)
+        return jsonify({"error": str(e), "message": "Terjadi kesalahan internal server."}), 500
 
 @app.route('/health', methods=['GET'])
 def health_check():
-    """Health check endpoint"""
-    return jsonify({
-        "status": "healthy",
-        "models_loaded": len(models),
-        "models_total": len(model_info),
-        "memory_usage": "CPU only"
-    })
+    """Endpoint untuk health check."""
+    return jsonify({"status": "healthy", "loaded_models_count": len(models), "message": "API berfungsi normal."})
 
+# --- Jalankan Aplikasi ---
 if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=7860, debug=False)
+    # Untuk Hugging Face Spaces, port biasanya 7860
+    # Menggunakan HOST dari environment variable jika tersedia, default ke 0.0.0.0
+    # Debug=False untuk produksi
+    app.run(host=os.getenv('HOST', '0.0.0.0'), port=int(os.getenv('PORT', 7860)), debug=False)