import gradio as gr
from transformers import pipeline

print("Loading Bangdim AI model...")

# Load model
chatbot = pipeline(
    "text-generation",
    model="microsoft/DialoGPT-small",
    device=-1  # Gunakan CPU
)

def respond(message, history):
    if not message:
        return "Halo kak! Ada yang bisa saya bantu? 😊"
    
    # Format prompt
    prompt = f"User: {message}\nAssistant:"
    
    # Generate response
    result = chatbot(
        prompt,
        max_length=100,
        temperature=0.8,
        do_sample=True,
        pad_token_id=50256
    )
    
    response = result[0]['generated_text']
    response = response.replace(prompt, "").strip()
    
    return response if response else "Maaf kak, saya kurang paham 😅"

# UI
with gr.Blocks() as demo:
    gr.Markdown("# 🤖 Bangdim AI Customer Service")
    gr.ChatInterface(respond)

demo.launch()