# Install paket yang dibutuhkan (jalankan ini di Colab atau lokal)
!pip install streamlit sentence-transformers transformers

import time
import streamlit as st
import torch
from sentence_transformers import SentenceTransformer, CrossEncoder

# Load Bi-Encoder dan Cross-Encoder
bi_encoder = SentenceTransformer('sentence-transformers/msmarco-distilbert-base-v3')
cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')

# UI dengan Streamlit
st.title("🔍 Perbandingan Bi-Encoder vs Cross-Encoder")
st.subheader("Masukkan Query dan Passages untuk melihat perbandingan ranking")

# Input untuk query dan passages
query = st.text_input("Masukkan Query:", "Apa manfaat AI dalam kehidupan sehari-hari?")
passages = st.text_area("Masukkan Passages (pisahkan dengan enter):", 
                        "AI membantu mengoptimalkan pencarian informasi dan otomatisasi tugas.\n"
                        "Teknologi AI telah berkembang pesat dalam bidang kesehatan dan pendidikan.\n"
                        "AI digunakan untuk meningkatkan efisiensi industri seperti manufaktur dan e-commerce.")

# Konversi input passages ke list
passages = passages.split("\n")

if st.button("Jalankan Model"):
    st.write("🚀 **Menganalisis passages dengan dua model...**")

    # --- Bi-Encoder ---
    start_time = time.time()
    query_emb = bi_encoder.encode(query)
    passage_embs = bi_encoder.encode(passages)
    scores_bi = [torch.cosine_similarity(torch.tensor(query_emb), torch.tensor(p_emb), dim=0).item() for p_emb in passage_embs]
    bi_time = time.time() - start_time  # Waktu eksekusi Bi-Encoder

    # --- Cross-Encoder ---
    start_time = time.time()
    scores_cross = cross_encoder.predict([[query, passage] for passage in passages])
    cross_time = time.time() - start_time  # Waktu eksekusi Cross-Encoder

    # Hitung Mean Reciprocal Rank (MRR)
    def compute_mrr(scores):
        ranked_scores = sorted(scores, reverse=True)
        if len(ranked_scores) == 0:
            return 0.0
        return 1 / (ranked_scores.index(max(scores)) + 1)

    mrr_bi = compute_mrr(scores_bi)
    mrr_cross = compute_mrr(scores_cross)

    # **Tampilkan Hasil**
    st.write("✅ **Hasil Ranking (Bi-Encoder)**")
    sorted_bi = sorted(zip(passages, scores_bi), key=lambda x: x[1], reverse=True)
    for i, (text, score) in enumerate(sorted_bi):
        st.write(f"{i+1}. **[{score:.4f}]** {text}")

    st.write("✅ **Hasil Ranking (Cross-Encoder)**")
    sorted_cross = sorted(zip(passages, scores_cross), key=lambda x: x[1], reverse=True)
    for i, (text, score) in enumerate(sorted_cross):
        st.write(f"{i+1}. **[{score:.4f}]** {text}")

    # **Tampilkan metrik**
    st.subheader("📊 Perbandingan Model")
    st.write(f"⏱ **Waktu Eksekusi Bi-Encoder:** {bi_time:.3f} detik")
    st.write(f"⏱ **Waktu Eksekusi Cross-Encoder:** {cross_time:.3f} detik")
    st.write(f"📈 **MRR Bi-Encoder:** {mrr_bi:.3f}")
    st.write(f"📈 **MRR Cross-Encoder:** {mrr_cross:.3f}")