123 lines
3.1 KiB
Python
123 lines
3.1 KiB
Python
|
|
import chromadb
|
|
from chromadb.utils import embedding_functions
|
|
|
|
|
|
# ==========================================
|
|
# PENGATURAN
|
|
# ==========================================
|
|
|
|
FOLDER_DATABASE = "./chroma_db"
|
|
|
|
MODEL_EMBEDDING = "paraphrase-multilingual-MiniLM-L12-v2"
|
|
|
|
|
|
# ==========================================
|
|
# MEMBUAT KONEKSI
|
|
# ==========================================
|
|
|
|
def hubungkan_database(ukuran):
|
|
client = chromadb.PersistentClient(
|
|
path=FOLDER_DATABASE
|
|
)
|
|
|
|
embedding = (
|
|
embedding_functions.SentenceTransformerEmbeddingFunction(
|
|
model_name=MODEL_EMBEDDING
|
|
)
|
|
)
|
|
|
|
nama_koleksi = f"dokumen_chunk_{ukuran}"
|
|
|
|
koleksi = client.get_collection(
|
|
name=nama_koleksi,
|
|
embedding_function=embedding
|
|
)
|
|
|
|
return koleksi
|
|
|
|
|
|
# ==========================================
|
|
# MELAKUKAN PENCARIAN
|
|
# ==========================================
|
|
|
|
def cari_dokumen():
|
|
print("\n===== SISTEM PENCARIAN DOKUMEN RAG =====")
|
|
print("Ukuran chunk tersedia: 500, 1000, 1500")
|
|
|
|
try:
|
|
ukuran = int(input("Pilih ukuran chunk: "))
|
|
|
|
if ukuran not in [500, 1000, 1500]:
|
|
print("Ukuran chunk tidak tersedia.")
|
|
return
|
|
|
|
koleksi = hubungkan_database(ukuran)
|
|
|
|
jumlah_data = koleksi.count()
|
|
|
|
if jumlah_data == 0:
|
|
print("Indeks masih kosong. Jalankan indeks.py dahulu.")
|
|
return
|
|
|
|
print(f"\nJumlah chunk dalam indeks: {jumlah_data}")
|
|
|
|
while True:
|
|
pertanyaan = input(
|
|
"\nMasukkan pertanyaan (ketik 'keluar' untuk berhenti): "
|
|
).strip()
|
|
|
|
if pertanyaan.lower() == "keluar":
|
|
break
|
|
|
|
if not pertanyaan:
|
|
print("Pertanyaan tidak boleh kosong.")
|
|
continue
|
|
|
|
jumlah_hasil = min(3, jumlah_data)
|
|
|
|
hasil = koleksi.query(
|
|
query_texts=[pertanyaan],
|
|
n_results=jumlah_hasil
|
|
)
|
|
|
|
print("\n===== HASIL PENCARIAN =====")
|
|
|
|
for nomor, (dokumen, metadata, jarak) in enumerate(
|
|
zip(
|
|
hasil["documents"][0],
|
|
hasil["metadatas"][0],
|
|
hasil["distances"][0]
|
|
),
|
|
start=1
|
|
):
|
|
print(f"\nHasil ke-{nomor}")
|
|
print(f"Sumber : {metadata['sumber']}")
|
|
print(f"Halaman: {metadata['halaman']}")
|
|
print(f"Chunk : {metadata['chunk']}")
|
|
print(f"Jarak : {jarak:.4f}")
|
|
print("Isi dokumen:")
|
|
print(dokumen)
|
|
print("--------------------------------")
|
|
|
|
print(
|
|
"Catatan: nilai jarak lebih kecil "
|
|
"berarti hasil lebih dekat."
|
|
)
|
|
|
|
|
|
except ValueError:
|
|
print("Masukkan angka ukuran chunk yang benar.")
|
|
|
|
except Exception as error:
|
|
print("Terjadi kesalahan:", error)
|
|
print("Pastikan indeks sudah dibuat dengan indeks.py.")
|
|
|
|
|
|
# ==========================================
|
|
# PROGRAM UTAMA
|
|
# ==========================================
|
|
|
|
if __name__ == "__main__":
|
|
cari_dokumen()
|