Spaces:

Lyon28
/

AI-Character-Chat

Sleeping

File size: 12,654 Bytes

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 typing import Optional

# Inisialisasi FastAPI
app = FastAPI(title="LyonPoy AI Chat")

# 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 ChatRequest(BaseModel):
    message: str
    model: Optional[str] = "gpt-2"

# Startup
@app.on_event("startup")
async def load_models():
    app.state.pipelines = {}
    os.environ['HF_HOME'] = '/tmp/.cache/huggingface'
    os.makedirs(os.environ['HF_HOME'], exist_ok=True)
    print("🤖 LyonPoy AI Chat Ready!")

# 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');

        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();
        }

        function showTyping() { typingIndicator.style.display = 'flex'; scrollToBottom(); }
        function hideTyping() { typingIndicator.style.display = 'none'; }

        async function sendMessage() {
            const message = chatInput.value.trim();
            if (!message) return;

            chatInput.disabled = true; sendButton.disabled = true;
            addMessage(message, true); chatInput.value = ''; showTyping();

            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();
        }

        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)

# 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
        if model_id not in app.state.pipelines:
            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
            )
        
        pipe = app.state.pipelines[model_id]
        
        # Process berdasarkan task
        if model_config["task"] == "text-generation":
            result = pipe(
                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"}
        
    except Exception as e:
        print(f"❌ Error: {e}")
        raise HTTPException(status_code=500, detail="Terjadi kesalahan")

# Health check
@app.get("/health")
async def health():
    return {"status": "healthy", "gpu": torch.cuda.is_available()}

# Run app
if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    uvicorn.run(app, host="0.0.0.0", port=port)