Update app.py

app.py

@@ -1,231 +1,319 @@
+import os
+import uvicorn
 from fastapi import FastAPI, HTTPException
+from fastapi.responses import HTMLResponse
+from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel
 from transformers import pipeline
 import torch
-from fastapi.middleware.cors import CORSMiddleware
-from typing import Dict, Any, Optional
-import os # Import os module
+from typing import Optional
 
-# Inisialisasi aplikasi FastAPI
-app = FastAPI(
-    title="LyonPoy Model Inference API",
-    description="API untuk mengakses 11 model machine learning",
-    version="1.0.0"
-)
+# Inisialisasi FastAPI
+app = FastAPI(title="LyonPoy AI Chat")
 
-# Konfigurasi CORS untuk frontend eksternal
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# Konfigurasi Model
-MODEL_MAP = {
-    "tinny-llama": "Lyon28/Tinny-Llama",
-    "pythia": "Lyon28/Pythia",
-    "bert-tinny": "Lyon28/Bert-Tinny",
-    "albert-base-v2": "Lyon28/Albert-Base-V2",
-    "t5-small": "Lyon28/T5-Small",
-    "gpt-2": "Lyon28/GPT-2",
-    "gpt-neo": "Lyon28/GPT-Neo",
-    "distilbert-base-uncased": "Lyon28/Distilbert-Base-Uncased",
-    "distil-gpt-2": "Lyon28/Distil_GPT-2",
-    "gpt-2-tinny": "Lyon28/GPT-2-Tinny",
-    "electra-small": "Lyon28/Electra-Small"
-}
-
-TASK_MAP = {
-    "text-generation": ["gpt-2", "gpt-neo", "distil-gpt-2", "gpt-2-tinny", "tinny-llama", "pythia"],
-    "text-classification": ["bert-tinny", "albert-base-v2", "distilbert-base-uncased", "electra-small"],
-    "text2text-generation": ["t5-small"]
+# All 11 models configuration
+MODELS = {
+    "tinny-llama": {
+        "name": "Tinny Llama",
+        "model_path": "Lyon28/Tinny-Llama",
+        "task": "text-generation"
+    },
+    "pythia": {
+        "name": "Pythia",
+        "model_path": "Lyon28/Pythia",
+        "task": "text-generation"
+    },
+    "bert-tinny": {
+        "name": "BERT Tinny",
+        "model_path": "Lyon28/Bert-Tinny",
+        "task": "text-classification"
+    },
+    "albert-base-v2": {
+        "name": "ALBERT Base V2",
+        "model_path": "Lyon28/Albert-Base-V2",
+        "task": "text-classification"
+    },
+    "t5-small": {
+        "name": "T5 Small",
+        "model_path": "Lyon28/T5-Small",
+        "task": "text2text-generation"
+    },
+    "gpt-2": {
+        "name": "GPT-2",
+        "model_path": "Lyon28/GPT-2",
+        "task": "text-generation"
+    },
+    "gpt-neo": {
+        "name": "GPT-Neo",
+        "model_path": "Lyon28/GPT-Neo",
+        "task": "text-generation"
+    },
+    "distilbert-base-uncased": {
+        "name": "DistilBERT",
+        "model_path": "Lyon28/Distilbert-Base-Uncased",
+        "task": "text-classification"
+    },
+    "distil-gpt-2": {
+        "name": "DistilGPT-2",
+        "model_path": "Lyon28/Distil_GPT-2",
+        "task": "text-generation"
+    },
+    "gpt-2-tinny": {
+        "name": "GPT-2 Tinny",
+        "model_path": "Lyon28/GPT-2-Tinny",
+        "task": "text-generation"
+    },
+    "electra-small": {
+        "name": "ELECTRA Small",
+        "model_path": "Lyon28/Electra-Small",
+        "task": "text-classification"
+    }
 }
 
-class InferenceRequest(BaseModel):
-    text: str
-    model_id: Optional[str] = "gpt-2"  # Default model
-    max_length: int = 100
-    temperature: float = 0.9
-    top_p: float = 0.95
-
-# Helper functions
-def get_task(model_id: str) -> str:
-    for task, models in TASK_MAP.items():
-        if model_id in models:
-            return task
-    # Default to text-generation if not found (or raise an error)
-    return "text-generation" 
+class ChatRequest(BaseModel):
+    message: str
+    model: Optional[str] = "gpt-2"
 
-# Event startup untuk inisialisasi model
+# Startup
 @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)
+    print("🤖 LyonPoy AI Chat Ready!")
 
-# Endpoint utama
-@app.get("/")
-async def root():
-    return {
-        "message": "Selamat datang di Lyon28 Model API",
-        "endpoints": {
-            "documentation": "/docs",
-            "model_list": "/models",
-            "health_check": "/health",
-            "inference_with_model": "/inference/{model_id}",
-            "inference_general": "/inference"
-        },
-        "total_models": len(MODEL_MAP),
-        "usage_examples": {
-            "specific_model": "POST /inference/gpt-2 with JSON body",
-            "general_inference": "POST /inference with model_id in JSON body"
+# Frontend route
+@app.get("/", response_class=HTMLResponse)
+async def get_frontend():
+    html_content = '''
+<!DOCTYPE html>
+<html lang="id">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>LyonPoy AI Chat</title>
+    <style>
+        * { margin: 0; padding: 0; box-sizing: border-box; }
+        body {
+            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
+            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+            height: 100vh; display: flex; justify-content: center; align-items: center;
         }
-    }
+        .chat-container {
+            width: 400px; height: 600px; background: #fff; border-radius: 15px;
+            box-shadow: 0 20px 40px rgba(0,0,0,0.15); display: flex; flex-direction: column; overflow: hidden;
+        }
+        .chat-header {
+            background: linear-gradient(135deg, #25d366, #128c7e); color: white;
+            padding: 20px; text-align: center;
+        }
+        .chat-header h1 { font-size: 18px; font-weight: 600; margin-bottom: 8px; }
+        .model-selector {
+            background: rgba(255,255,255,0.2); border: none; color: white;
+            padding: 8px 12px; border-radius: 20px; font-size: 12px; cursor: pointer;
+        }
+        .chat-messages {
+            flex: 1; padding: 20px; overflow-y: auto; background: #f0f0f0;
+            display: flex; flex-direction: column; gap: 15px;
+        }
+        .message {
+            max-width: 80%; padding: 12px 16px; border-radius: 15px;
+            font-size: 14px; line-height: 1.4; animation: slideIn 0.3s ease;
+        }
+        .message.user {
+            background: #25d366; color: white; align-self: flex-end; border-bottom-right-radius: 5px;
+        }
+        .message.bot {
+            background: white; color: #333; align-self: flex-start;
+            border-bottom-left-radius: 5px; box-shadow: 0 2px 5px rgba(0,0,0,0.1);
+        }
+        .message-time { font-size: 11px; opacity: 0.7; margin-top: 5px; }
+        .chat-input-container {
+            padding: 20px; background: white; border-top: 1px solid #e0e0e0;
+            display: flex; gap: 10px; align-items: center;
+        }
+        .chat-input {
+            flex: 1; padding: 12px 16px; border: 1px solid #e0e0e0;
+            border-radius: 25px; font-size: 14px; outline: none;
+        }
+        .chat-input:focus { border-color: #25d366; box-shadow: 0 0 0 2px rgba(37, 211, 102, 0.2); }
+        .send-button {
+            background: #25d366; color: white; border: none; border-radius: 50%;
+            width: 45px; height: 45px; cursor: pointer; display: flex;
+            align-items: center; justify-content: center;
+        }
+        .send-button:hover { background: #128c7e; }
+        .send-button:disabled { background: #ccc; cursor: not-allowed; }
+        .welcome-message {
+            text-align: center; color: #666; font-size: 13px;
+            padding: 20px; border-radius: 10px; background: rgba(255,255,255,0.7);
+        }
+        .typing-indicator {
+            display: none; align-items: center; gap: 5px; padding: 12px 16px;
+            background: white; border-radius: 15px; align-self: flex-start;
+        }
+        .typing-dot {
+            width: 8px; height: 8px; background: #999; border-radius: 50%;
+            animation: typing 1.4s infinite;
+        }
+        .typing-dot:nth-child(2) { animation-delay: 0.2s; }
+        .typing-dot:nth-child(3) { animation-delay: 0.4s; }
+        @keyframes typing { 0%, 60%, 100% { transform: translateY(0); } 30% { transform: translateY(-10px); } }
+        @keyframes slideIn { from { opacity: 0; transform: translateY(20px); } to { opacity: 1; transform: translateY(0); } }
+        @media (max-width: 480px) { .chat-container { width: 100vw; height: 100vh; border-radius: 0; } }
+    </style>
+</head>
+<body>
+    <div class="chat-container">
+        <div class="chat-header">
+            <h1>🤖 LyonPoy AI Chat</h1>
+            <select class="model-selector" id="modelSelect">
+                <option value="gpt-2">GPT-2 (General)</option>
+                <option value="tinny-llama">Tinny Llama</option>
+                <option value="pythia">Pythia</option>
+                <option value="gpt-neo">GPT-Neo</option>
+                <option value="distil-gpt-2">DistilGPT-2</option>
+                <option value="gpt-2-tinny">GPT-2 Tinny</option>
+                <option value="bert-tinny">BERT Tinny</option>
+                <option value="albert-base-v2">ALBERT Base V2</option>
+                <option value="distilbert-base-uncased">DistilBERT</option>
+                <option value="electra-small">ELECTRA Small</option>
+                <option value="t5-small">T5 Small</option>
+            </select>
+        </div>
+        <div class="chat-messages" id="chatMessages">
+            <div class="welcome-message">
+                👋 Halo! Saya LyonPoy AI Assistant.<br>
+                Pilih model di atas dan mulai chat dengan saya!
+            </div>
+        </div>
+        <div class="typing-indicator" id="typingIndicator">
+            <div class="typing-dot"></div><div class="typing-dot"></div><div class="typing-dot"></div>
+        </div>
+        <div class="chat-input-container">
+            <input type="text" class="chat-input" id="chatInput" placeholder="Ketik pesan..." maxlength="500">
+            <button class="send-button" id="sendButton">➤</button>
+        </div>
+    </div>
+    <script>
+        const chatMessages = document.getElementById('chatMessages');
+        const chatInput = document.getElementById('chatInput');
+        const sendButton = document.getElementById('sendButton');
+        const modelSelect = document.getElementById('modelSelect');
+        const typingIndicator = document.getElementById('typingIndicator');
 
-# Endpoint untuk list model
-@app.get("/models")
-async def list_models():
-    return {
-        "available_models": list(MODEL_MAP.keys()),
-        "total_models": len(MODEL_MAP)
-    }
+        function scrollToBottom() { chatMessages.scrollTop = chatMessages.scrollHeight; }
+        
+        function addMessage(content, isUser = false) {
+            const messageDiv = document.createElement('div');
+            messageDiv.className = `message ${isUser ? 'user' : 'bot'}`;
+            const time = new Date().toLocaleTimeString('id-ID', { hour: '2-digit', minute: '2-digit' });
+            messageDiv.innerHTML = `${content}<div class="message-time">${time}</div>`;
+            chatMessages.appendChild(messageDiv);
+            scrollToBottom();
+        }
 
-# Endpoint health check
-@app.get("/health")
-async def health_check():
-    return {
-        "status": "healthy",
-        "gpu_available": torch.cuda.is_available(),
-        "gpu_type": torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU-only"
-    }
+        function showTyping() { typingIndicator.style.display = 'flex'; scrollToBottom(); }
+        function hideTyping() { typingIndicator.style.display = 'none'; }
 
-# NEW: General inference endpoint (handles POST /inference)
-@app.post("/inference")
-async def general_inference(request: InferenceRequest):
-    """
-    General inference endpoint that accepts model_id in the request body
-    """
-    return await process_inference(request.model_id, request)
+        async function sendMessage() {
+            const message = chatInput.value.trim();
+            if (!message) return;
 
-# Endpoint inference dengan model_id di path
-@app.post("/inference/{model_id}")
-async def model_inference(model_id: str, request: InferenceRequest):
-    """
-    Specific model inference endpoint with model_id in path
-    """
-    return await process_inference(model_id, request)
+            chatInput.disabled = true; sendButton.disabled = true;
+            addMessage(message, true); chatInput.value = ''; showTyping();
 
-# Shared inference processing function
-async def process_inference(model_id: str, request: InferenceRequest):
-    try:
-        # Pastikan model_id dalam lowercase agar sesuai dengan MODEL_MAP
-        model_id = model_id.lower() 
+            try {
+                const response = await fetch('/chat', {
+                    method: 'POST',
+                    headers: { 'Content-Type': 'application/json' },
+                    body: JSON.stringify({ message: message, model: modelSelect.value })
+                });
+                const data = await response.json();
+                hideTyping();
+                if (data.status === 'success') {
+                    addMessage(data.response);
+                } else {
+                    addMessage('❌ Maaf, terjadi kesalahan. Coba lagi nanti.');
+                }
+            } catch (error) {
+                hideTyping();
+                addMessage('❌ Tidak dapat terhubung ke server.');
+            }
+            chatInput.disabled = false; sendButton.disabled = false; chatInput.focus();
+        }
 
-        # Validasi model ID
-        if model_id not in MODEL_MAP:
-            available_models = ", ".join(MODEL_MAP.keys())
-            raise HTTPException(
-                status_code=404,
-                detail=f"Model '{model_id}' tidak ditemukan. Model yang tersedia: {available_models}"
-            )
+        sendButton.addEventListener('click', sendMessage);
+        chatInput.addEventListener('keypress', (e) => { if (e.key === 'Enter') sendMessage(); });
+        modelSelect.addEventListener('change', () => {
+            const modelName = modelSelect.options[modelSelect.selectedIndex].text;
+            addMessage(`🔄 Model diubah ke: ${modelName}`);
+        });
+        window.addEventListener('load', () => chatInput.focus());
+    </script>
+</body>
+</html>
+    '''
+    return HTMLResponse(content=html_content)
 
-        # Dapatkan task yang sesuai
-        task = get_task(model_id)
+# Chat API
+@app.post("/chat")
+async def chat(request: ChatRequest):
+    try:
+        model_id = request.model.lower()
+        if model_id not in MODELS:
+            raise HTTPException(status_code=400, detail="Model tidak tersedia")
+        
+        model_config = MODELS[model_id]
         
-        # Load model jika belum ada di memory
+        # Load model jika belum ada
         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
-
-            try:
-                app.state.pipelines[model_id] = pipeline(
-                    task=task,
-                    model=MODEL_MAP[model_id],
-                    device=device_to_use,
-                    torch_dtype=dtype_to_use
-                )
-                print(f"✅ Model {model_id} berhasil dimuat!")
-            except Exception as load_error:
-                print(f"❌ Gagal memuat model {model_id}: {load_error}")
-                raise HTTPException(
-                    status_code=503,
-                    detail=f"Gagal memuat model {model_id}. Coba lagi nanti."
-                )
-
-        pipe = app.state.pipelines[model_id]
-
-        # Proses berdasarkan task
-        if task == "text-generation":
-            result = pipe(
-                request.text,
-                max_length=request.max_length,
-                temperature=request.temperature,
-                top_p=request.top_p,
-                do_sample=True
-            )[0]['generated_text']
+            print(f"⏳ Loading {model_config['name']}...")
+            device = 0 if torch.cuda.is_available() else -1
+            dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+            
+            app.state.pipelines[model_id] = pipeline(
+                task=model_config["task"],
+                model=model_config["model_path"],
+                device=device,
+                torch_dtype=dtype
+            )
         
-        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'],
-                "confidence": round(output['score'], 4)
-            }
+        pipe = app.state.pipelines[model_id]
         
-        elif task == "text2text-generation":
-            # Untuk text2text-generation, output juga list of dict
+        # Process berdasarkan task
+        if model_config["task"] == "text-generation":
             result = pipe(
-                request.text,
-                max_length=request.max_length
+                request.message,
+                max_length=min(len(request.message.split()) + 50, 200),
+                temperature=0.7,
+                do_sample=True,
+                pad_token_id=pipe.tokenizer.eos_token_id
             )[0]['generated_text']
+            
+            # Clean output
+            if result.startswith(request.message):
+                result = result[len(request.message):].strip()
+            
+        elif model_config["task"] == "text-classification":
+            output = pipe(request.message)[0]
+            result = f"Sentimen: {output['label']} (Confidence: {output['score']:.2f})"
+            
+        elif model_config["task"] == "text2text-generation":
+            result = pipe(request.message, max_length=150)[0]['generated_text']
+        
+        return {"response": result, "model": model_config["name"], "status": "success"}
         
-        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,
-            "model_used": model_id,
-            "task": task,
-            "status": "success"
-        }
-    
-    except HTTPException as he:
-        # Re-raise HTTP exceptions
-        raise he
     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
+        print(f"❌ Error: {e}")
+        raise HTTPException(status_code=500, detail="Terjadi kesalahan")
 
-        raise HTTPException(
-            status_code=500,
-            detail=f"Error processing request: {str(e)}. Cek log server untuk detail."
-        )
+# Health check
+@app.get("/health")
+async def health():
+    return {"status": "healthy", "gpu": torch.cuda.is_available()}
 
-# Error handler untuk 404
-@app.exception_handler(404)
-async def not_found_handler(request, exc):
-    return {
-        "error": "Endpoint tidak ditemukan",
-        "available_endpoints": [
-            "GET /",
-            "GET /models", 
-            "GET /health",
-            "POST /inference",
-            "POST /inference/{model_id}"
-        ],
-        "tip": "Gunakan /docs untuk dokumentasi lengkap"
-    }
+# Run app
+if __name__ == "__main__":
+    port = int(os.environ.get("PORT", 7860))
+    uvicorn.run(app, host="0.0.0.0", port=port)