tugas-josjis/lab3_hit_at_k.py

from pathlib import Path
from pypdf import PdfReader
import chromadb
from chromadb.utils import embedding_functions

# ==========================================
# 1. MEMBACA DOKUMEN PDF
# ==========================================

def baca_pdf(path):
    reader = PdfReader(path)
    return "\n".join((p.extract_text() or "") for p in reader.pages)


# ==========================================
# 2. MEMBAGI DOKUMEN MENJADI CHUNK
# ==========================================

def potong(teks, ukuran=800, tumpang=100):
    hasil = []
    mulai = 0

    while mulai < len(teks):
        hasil.append(teks[mulai:mulai + ukuran])
        mulai += ukuran - tumpang

    return [h for h in hasil if h.strip()]


# ==========================================
# 3. PENGATURAN EMBEDDING DAN DATABASE
# ==========================================

emb = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name="paraphrase-multilingual-MiniLM-L12-v2"
)

client = chromadb.PersistentClient(path="./chroma_db")


# ==========================================
# 4. MEMBUAT KOLEKSI
# ==========================================

def buat_koleksi(nama, ukuran, tumpang):

    try:
        client.delete_collection(nama)
    except Exception:
        pass

    kol = client.get_or_create_collection(
        name=nama,
        embedding_function=emb,
        metadata={"hnsw:space": "cosine"}
    )

    for pdf in Path("dokumen").glob("*.pdf"):

        teks = baca_pdf(pdf)
        chunks = potong(teks, ukuran, tumpang)

        if not chunks:
            continue

        kol.upsert(
            ids=[
                f"{nama}-{pdf.stem}-{i}"
                for i in range(len(chunks))
            ],
            documents=chunks,
            metadatas=[
                {"sumber": pdf.name, "chunk": i}
                for i in range(len(chunks))
            ]
        )

    return kol


# ==========================================
# 5. 15 PERTANYAAN DAN KATA KUNCI
# ==========================================

UJI = [
    ("What is a sample space?", "sample space"),
    ("What is classical probability?", "classical probability"),
    ("What is empirical probability?", "empirical probability"),
    ("What is subjective probability?", "subjective probability"),
    ("What is a simple event?", "simple event"),
    ("What is the multiplication rule?", "multiplication rule"),
    ("What is the addition rule?", "addition rule"),
    ("What are mutually exclusive events?", "mutually exclusive"),
    ("What is conditional probability?", "conditional probability"),
    ("What is the law of large numbers?", "law of large numbers"),
    ("What is a discrete probability distribution?", "discrete probability distribution"),
    ("What is a binomial experiment?", "binomial experiment"),
    ("What is the fundamental counting principle?", "fundamental counting principle"),
    ("What is the difference between permutations and combinations?", "permutation"),
    ("What is a random variable?", "random variable"),
]


# ==========================================
# 6. FUNGSI HIT@K
# ==========================================

def hit_at_k(koleksi, uji, k):

    benar = 0
    jumlah = min(k, koleksi.count())

    if jumlah == 0:
        return 0

    for pertanyaan, kunci in uji:

        docs = koleksi.query(
            query_texts=[pertanyaan],
            n_results=jumlah
        )["documents"][0]

        if any(kunci.lower() in d.lower() for d in docs):
            benar += 1

    return benar / len(uji)


# ==========================================
# 7. MEMBUAT DUA KOLEKSI BERBEDA
# ==========================================

kol_500 = buat_koleksi("chunk_500", 500, 75)
kol_1500 = buat_koleksi("chunk_1500", 1500, 225)


# ==========================================
# 8. MENAMPILKAN TABEL HASIL HIT@K
# ==========================================

print("\n========== TABEL HIT@K ==========")

print(f"{'k':<5}{'Chunk 500':<15}{'Chunk 1500':<15}")

hasil = {}

for k in (1, 3, 5):

    h500 = hit_at_k(kol_500, UJI, k)
    h1500 = hit_at_k(kol_1500, UJI, k)

    hasil[k] = (h500, h1500)

    print(f"{k:<5}{h500:<15.2f}{h1500:<15.2f}")


# ==========================================
# 9. MENAMPILKAN PERTANYAAN YANG GAGAL
# ==========================================

print("\n========== PERTANYAAN GAGAL DI k=3 ==========")

gagal = []

for pertanyaan, kunci in UJI:

    docs_500 = kol_500.query(
        query_texts=[pertanyaan],
        n_results=min(3, kol_500.count())
    )["documents"][0]

    docs_1500 = kol_1500.query(
        query_texts=[pertanyaan],
        n_results=min(3, kol_1500.count())
    )["documents"][0]

    ok_500 = any(kunci.lower() in d.lower() for d in docs_500)
    ok_1500 = any(kunci.lower() in d.lower() for d in docs_1500)

    if not ok_500 or not ok_1500:

        gagal.append(pertanyaan)

        print(f"\nPertanyaan: {pertanyaan}")
        print(f"Kata kunci: {kunci}")
        print(f"Chunk 500 : {'Berhasil' if ok_500 else 'Gagal'}")
        print(f"Chunk 1500: {'Berhasil' if ok_1500 else 'Gagal'}")

if not gagal:
    print("Semua pertanyaan berhasil ditemukan pada kedua koleksi.")


# ==========================================
# 10. KESIMPULAN OTOMATIS
# ==========================================

print("\n========== KESIMPULAN ==========")

rata_500 = sum(v[0] for v in hasil.values()) / len(hasil)
rata_1500 = sum(v[1] for v in hasil.values()) / len(hasil)

print(f"Rata-rata Hit@K Chunk 500  : {rata_500:.2f}")
print(f"Rata-rata Hit@K Chunk 1500 : {rata_1500:.2f}")

if rata_500 > rata_1500:
    print("Chunk 500 memiliki hasil pencarian lebih baik.")
elif rata_1500 > rata_500:
    print("Chunk 1500 memiliki hasil pencarian lebih baik.")
else:
    print("Kedua ukuran chunk memiliki hasil rata-rata yang sama.")

print(
    f"Jumlah pertanyaan yang gagal pada minimal satu koleksi "
    f"di k=3: {len(gagal)} dari {len(UJI)}."
)

print(
    "Kegagalan pencarian dapat disebabkan oleh perbedaan istilah, "
    "informasi yang tersebar di beberapa chunk, atau hasil embedding "
    "yang kurang relevan."
)

print("\n========== PENGUJIAN SELESAI ==========")