import gradio as gr
from transformers import pipeline
import torch

# Muat model mt5-small yang andal dan multilingual
print("Memuat pipeline untuk model 'google/mt5-small'...")
chatbot_pipeline = pipeline(
    "text2text-generation",
    model="google/mt5-small"
)
print("Pipeline model berhasil dimuat.")

# Fungsi untuk memproses input dan memberikan output
def chat(prompt):
    print(f"Menerima prompt: {prompt}")
    
    # --- PERUBAHAN UTAMA DI SINI ---
    # Kita ganti dari beam search ke sampling untuk hasil yang lebih beragam
    response = chatbot_pipeline(
        prompt,
        max_length=128,
        do_sample=True,
        temperature=0.8,
        top_k=50,
        top_p=0.95
    )
    
    generated_text = response[0]['generated_text']
    
    print(f"Jawaban yang dihasilkan: {generated_text}")
    return generated_text

# Buat antarmuka web dengan Gradio
web_interface = gr.Interface(
    fn=chat,
    inputs=gr.Textbox(lines=5, label="Input Pengguna", placeholder="Ketik pertanyaan atau kalimat Anda di sini..."),
    outputs=gr.Textbox(label="Jawaban Chatbot"),
    title="🇮🇩 Chatbot Indonesia (Google MT5)",
    description="Chatbot ini ditenagai oleh model 'google/mt5-small', sebuah model multilingual yang andal dari Google.",
    examples=[
        ["Siapa saja pahlawan nasional dari Jawa Barat?"],
        ["Jelaskan cara kerja internet dalam tiga kalimat"],
        ["Buatlah kalimat yang mengandung kata 'ekonomi' dan 'digital'"]
    ],
    allow_flagging="never"
)

# Jalankan aplikasi dengan konfigurasi server yang benar
if __name__ == "__main__":
    web_interface.launch(server_name="0.0.0.0", server_port=7860)