import gradio as gr
from transformers import pipeline
import torch

# Cek GPU
device = 0 if torch.cuda.is_available() else -1

# Load model (akan otomatis download pertama kali)
print("Loading AI model...")
generator = pipeline(
    'text-generation',
    model='microsoft/DialoGPT-small',  # Model kecil, cepat
    device=device
)

def chat(message, history):
    # Build context
    context = ""
    for h in history[-3:]:
        context += f"User: {h[0]}\nAI: {h[1]}\n"
    
    prompt = f"{context}User: {message}\nAI:"
    
    # Generate response
    result = generator(
        prompt,
        max_length=150,
        temperature=0.8,
        do_sample=True,
        pad_token_id=50256
    )
    
    response = result[0]['generated_text']
    response = response.replace(prompt, "").strip()
    
    if not response:
        response = "Maaf kak, saya kurang paham. Bisa diulang? 😊"
    
    history.append((message, response))
    return history, history

# UI
with gr.Blocks() as demo:
    gr.Markdown("# 🤖 Bangdim AI")
    chatbot = gr.Chatbot()
    msg = gr.Textbox(label="Pesan")
    clear = gr.Button("Clear")
    
    msg.submit(chat, [msg, chatbot], [chatbot, msg])
    clear.click(lambda: None, None, chatbot)

demo.launch()