from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import gradio as gr
import numpy as np # Import numpy jika belum ada

# --- 1. Inisialisasi Model dan Tokenizer (Dilakukan Sekali Saat Aplikasi Dimulai) ---
# Pastikan 'model_name' ini adalah model yang sudah kamu unggah ke Hugging Face Hub
# atau model publik lain yang ingin kamu gunakan.
model_name = "atsnetwork/my-custom-tinyllama-chatbot"
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token # Pastikan token padding diatur
tokenizer.padding_side = "right" # Penting untuk efisiensi saat memproses sequence

# --- 2. Fungsi Pemroses Utama untuk Chatbot (Dipanggil Oleh Gradio) ---
# Fungsi ini mengintegrasikan logika inferensi dan analisis developer info
def generate_response_with_dev_info(prompt, max_new_tokens=100, temperature=0.6, top_k=30):
    formatted_prompt = f"<s>[INST] {prompt} [/INST]"
    inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)

    # Generate dengan output_scores=True untuk analisis probabilitas
    outputs = model.generate(
        **inputs,
        max_new_tokens=max_new_tokens,
        do_sample=True,
        temperature=temperature,
        top_k=top_k,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.pad_token_id,
        return_dict_in_generate=True,
        output_scores=True
    )

    generated_ids = outputs.sequences[0]
    generated_text = tokenizer.decode(generated_ids, skip_special_tokens=True)

    # Ekstrak jawaban bersih
    answer = ""
    start_answer = generated_text.find("[/INST]")
    if start_answer != -1:
        answer = generated_text[start_answer + len("[/INST]"):].strip()
        if answer.endswith("</s>"):
            answer = answer[:-len("</s>")].strip()
    else:
        answer = generated_text.strip() # Fallback jika format tidak ditemukan

    # --- Analisis Konfidensi Berbasis ENTROPY (PENGGANTI avg_max_prob) ---
    avg_entropy = 0
    total_generated_tokens = 0
    if outputs.scores:
        for score_tensor in outputs.scores:
            probabilities = torch.softmax(score_tensor, dim=-1)
            # Menghitung entropy untuk setiap distribusi probabilitas
            epsilon = 1e-9 # Tambahkan sedikit epsilon untuk menghindari log(0)
            entropy = -torch.sum(probabilities * torch.log(probabilities + epsilon), dim=-1).item()
            avg_entropy += entropy
            total_generated_tokens += 1
        
        if total_generated_tokens > 0:
            avg_entropy /= total_generated_tokens
        else:
            # Jika tidak ada token yang dihasilkan setelah prompt (kasus jarang)
            avg_entropy = float('inf') 

    # Batas ambang keyakinan (contoh) untuk ENTROPY
    # Penting: Nilai ini HARUS ditentukan setelah analisis empiris pada modelmu.
    # Nilai entropy yang lebih RENDAH berarti model lebih YAKIN.
    # Nilai entropy yang lebih TINGGI berarti model lebih TIDAK YAKIN.
    # Sebagai titik awal, kamu bisa coba sekitar 1.0 atau lebih,
    # tetapi validasi empiris sangat dianjurkan.
    entropy_threshold = 1.0 # <--- SESUAIKAN NILAI INI BERDASARKAN ANALISISMU!

    if avg_entropy > entropy_threshold: # <--- Perhatikan operatornya berubah (lebih besar = LOW_CONFIDENCE)
        confidence_status = "LOW_CONFIDENCE: Model mungkin tidak memiliki pola yang jelas (Entropy Tinggi)."
    else:
        confidence_status = "HIGH_CONFIDENCE"

    # --- Analisis Frasa "Tidak Tahu" yang Dilatih ---
    is_explicitly_unknown = False
    explicit_unknown_reason = ""
    unknown_phrases = [
        "maaf, saya tidak memiliki informasi",
        "saya tidak familiar dengan",
        "di luar cakupan data pelatihan saya",
        "saya tidak tahu",
        "tidak dapat menemukan informasi"
    ]
    answer_lower = answer.lower()
    for phrase in unknown_phrases:
        if phrase in answer_lower:
            is_explicitly_unknown = True
            explicit_unknown_reason = f"Model menggunakan frasa 'tidak tahu': '{phrase}'"
            break # Hentikan setelah menemukan frasa pertama

    # --- Gabungkan Informasi untuk Developer ---
    developer_info = {
        "confidence_score": f"{avg_entropy:.4f}", # Sekarang ini adalah entropy
        "confidence_status": confidence_status,
        "explicit_unknown_phrase_detected": is_explicitly_unknown,
        "explicit_unknown_reason": explicit_unknown_reason if is_explicitly_unknown else "Tidak ada frasa 'tidak tahu' eksplisit.",
        # "raw_generated_text": generated_text # Bisa diaktifkan untuk debug, tapi akan terlihat di UI
    }

    # Untuk tampilan user, hanya tampilkan jawabannya.
    # Informasi developer bisa ditampilkan di antarmuka terpisah atau log.
    return answer, developer_info

# --- 3. Fungsi Adaptor untuk Gradio Interface (Mengonversi Dictionary info menjadi String) ---
# Gradio Interface mengharapkan output string/angka, bukan dictionary.
# Fungsi ini akan mengubah dictionary developer_info menjadi string yang mudah dibaca.
def gradio_interface_fn(prompt):
    answer, dev_info = generate_response_with_dev_info(prompt)
    
    # Format developer info untuk ditampilkan di Gradio
    dev_info_str = "--- Developer Info ---\n"
    dev_info_str += f"Confidence Score (Entropy): {dev_info['confidence_score']} ({dev_info['confidence_status']})\n"
    dev_info_str += f"Explicit Unknown Phrase Detected: {dev_info['explicit_unknown_phrase_detected']}\n"
    dev_info_str += f"Reason: {dev_info['explicit_unknown_reason']}\n"
    # dev_info_str += f"Raw Generated Text: {dev_info['raw_generated_text']}\n" 

    return answer, dev_info_str

# --- 4. Inisialisasi Antarmuka Gradio ---
# Ini yang akan membangun UI di Hugging Face Space.
iface = gr.Interface(
    fn=gradio_interface_fn, # Fungsi yang akan dipanggil saat ada input
    inputs=gr.Textbox(lines=2, label="Your Question"), # Input berupa kotak teks untuk pertanyaan
    outputs=[
        gr.Textbox(label="Chatbot Response", lines=5), # Output pertama untuk jawaban chatbot
        gr.Textbox(label="Developer Information", lines=5) # Output kedua untuk informasi developer
    ],
    title="TinyLlama Custom Chatbot with Developer Insights 🚀",
    description="Ask anything and get a response from the chatbot. Additional information for developers will be displayed below."
)

# --- 5. Jalankan Aplikasi Gradio ---
iface.launch()