import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import os # <-- Tambahkan ini

# --- 1. Muat Model dan Tokenizer dari Hugging Face Hub ---
# Ganti dengan username dan nama repo Anda yang sudah di-fine-tune
MODEL_REPO = "Wiefdw/modelAnevia-v2"

# Ambil token dari Space Secrets (lebih aman!)
# Pastikan Anda sudah mengatur secret bernama "HF_TOKEN" di pengaturan Space Anda
HF_TOKEN = os.getenv("HF_TOKEN") 

# Cek apakah token ditemukan
if not HF_TOKEN:
    raise ValueError("Hugging Face token tidak ditemukan. Mohon atur 'HF_TOKEN' di Space Secrets.")

print(f"Memuat model dari: {MODEL_REPO}")
try:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO, token=HF_TOKEN)
    model = AutoModelForCausalLM.from_pretrained(MODEL_REPO, token=HF_TOKEN)
    print("Model berhasil dimuat.")
except Exception as e:
    print(f"Gagal memuat model. Pastikan nama repo dan token sudah benar. Error: {e}")
    # Hentikan aplikasi jika model gagal dimuat
    raise

# Pindahkan model ke GPU jika tersedia, untuk inferensi yang lebih cepat
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
print(f"Model berjalan di: {device}")


# --- 2. Fungsi untuk Menghasilkan Respon Chatbot ---
def generate_response(message, chat_history):
    """
    Fungsi ini mengambil pesan pengguna dan riwayat obrolan, 
    lalu menghasilkan respons dari model.
    """
    # Format ulang prompt dengan riwayat obrolan
    prompt_history = ""
    for user_input, bot_response in chat_history:
        prompt_history += f"{user_input}{tokenizer.eos_token}{bot_response}{tokenizer.eos_token}"
    
    # Gabungkan dengan pesan baru
    prompt = f"{prompt_history}{message}{tokenizer.eos_token}"
    
    # Tokenisasi prompt
    inputs = tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True).to(device)
    
    # Hasilkan respons dari model
    outputs = model.generate(
        **inputs,
        max_new_tokens=100,
        pad_token_id=tokenizer.eos_token_id,
        no_repeat_ngram_size=3,
        do_sample=True,
        top_k=50,
        top_p=0.95,
        temperature=0.7
    )
    
    # Decode hanya token yang baru dihasilkan
    new_tokens = outputs[0, inputs['input_ids'].shape[-1]:]
    bot_response = tokenizer.decode(new_tokens, skip_special_tokens=True)
    
    # Tambahkan penanganan jika respons kosong
    return bot_response if bot_response else "Maaf, saya tidak bisa merespons saat ini."

# --- 3. Buat Antarmuka Gradio ---
chatbot_interface = gr.ChatInterface(
    fn=generate_response,
    title="🩺 Chatbot Informasi Anemia",
    description="Tanyakan apa saja tentang anemia. Model ini di-fine-tune dari DialoGPT-medium.",
    examples=[
        ["Apa itu anemia?"],
        ["Makanan apa yang baik untuk penderita anemia?"],
        ["Apakah anemia bisa sembuh?"]
    ],
    theme="soft",
    retry_btn=None,
    undo_btn="Hapus Pesan Terakhir",
    clear_btn="Bersihkan Obrolan",
)

# --- 4. Jalankan Aplikasi ---
if __name__ == "__main__":
    chatbot_interface.launch()