175 lines
4.3 KiB
Python
175 lines
4.3 KiB
Python
|
|
from pathlib import Path
|
|
from pypdf import PdfReader
|
|
import chromadb
|
|
from chromadb.utils import embedding_functions
|
|
|
|
|
|
# ==========================================
|
|
# PENGATURAN
|
|
# ==========================================
|
|
|
|
FOLDER_DOKUMEN = Path("dokumen")
|
|
FOLDER_DATABASE = "./chroma_db"
|
|
|
|
UKURAN_CHUNK = [500, 1000, 1500]
|
|
MODEL_EMBEDDING = "paraphrase-multilingual-MiniLM-L12-v2"
|
|
|
|
|
|
# ==========================================
|
|
# MEMBACA FILE PDF
|
|
# ==========================================
|
|
|
|
def baca_pdf(path):
|
|
reader = PdfReader(path)
|
|
daftar_halaman = []
|
|
|
|
for nomor, halaman in enumerate(reader.pages, start=1):
|
|
teks = halaman.extract_text() or ""
|
|
|
|
if teks.strip():
|
|
daftar_halaman.append({
|
|
"teks": teks,
|
|
"halaman": nomor
|
|
})
|
|
|
|
return daftar_halaman
|
|
|
|
|
|
# ==========================================
|
|
# MEMOTONG TEKS MENJADI CHUNK
|
|
# ==========================================
|
|
|
|
def potong(teks, ukuran=500, tumpang=50):
|
|
hasil = []
|
|
mulai = 0
|
|
|
|
while mulai < len(teks):
|
|
bagian = teks[mulai:mulai + ukuran]
|
|
|
|
if bagian.strip():
|
|
hasil.append(bagian)
|
|
|
|
mulai += ukuran - tumpang
|
|
|
|
return hasil
|
|
|
|
|
|
# ==========================================
|
|
# MEMBUAT INDEKS CHROMADB
|
|
# ==========================================
|
|
|
|
def buat_indeks():
|
|
if not FOLDER_DOKUMEN.exists():
|
|
print("Folder dokumen tidak ditemukan.")
|
|
return
|
|
|
|
daftar_pdf = list(FOLDER_DOKUMEN.glob("*.pdf"))
|
|
|
|
if not daftar_pdf:
|
|
print("Belum ada file PDF di folder dokumen.")
|
|
return
|
|
|
|
print("Menghubungkan ke ChromaDB...")
|
|
|
|
client = chromadb.PersistentClient(
|
|
path=FOLDER_DATABASE
|
|
)
|
|
|
|
embedding = (
|
|
embedding_functions.SentenceTransformerEmbeddingFunction(
|
|
model_name=MODEL_EMBEDDING
|
|
)
|
|
)
|
|
|
|
print("Model embedding siap digunakan.")
|
|
|
|
# Membuat koleksi untuk setiap ukuran chunk.
|
|
for ukuran in UKURAN_CHUNK:
|
|
|
|
nama_koleksi = f"dokumen_chunk_{ukuran}"
|
|
|
|
koleksi = client.get_or_create_collection(
|
|
name=nama_koleksi,
|
|
embedding_function=embedding,
|
|
metadata={"hnsw:space": "cosine"}
|
|
)
|
|
|
|
print("\n================================")
|
|
print(f"MEMBUAT INDEKS CHUNK {ukuran}")
|
|
print("================================")
|
|
|
|
total_chunk = 0
|
|
|
|
for pdf in daftar_pdf:
|
|
halaman_pdf = baca_pdf(pdf)
|
|
|
|
ids = []
|
|
documents = []
|
|
metadatas = []
|
|
|
|
for data in halaman_pdf:
|
|
teks = data["teks"]
|
|
|
|
overlap = int(ukuran * 0.10)
|
|
|
|
chunks = potong(
|
|
teks,
|
|
ukuran=ukuran,
|
|
tumpang=overlap
|
|
)
|
|
|
|
for nomor, chunk in enumerate(chunks):
|
|
id_chunk = (
|
|
f"{pdf.stem}_halaman_"
|
|
f"{data['halaman']}_chunk_"
|
|
f"{nomor}_ukuran_{ukuran}"
|
|
)
|
|
|
|
ids.append(id_chunk)
|
|
documents.append(chunk)
|
|
|
|
metadatas.append({
|
|
"sumber": pdf.name,
|
|
"halaman": data["halaman"],
|
|
"chunk": nomor,
|
|
"ukuran_chunk": ukuran
|
|
})
|
|
|
|
if ids:
|
|
# Simpan dalam kelompok agar lebih aman
|
|
# untuk dokumen dengan banyak chunk.
|
|
for awal in range(0, len(ids), 100):
|
|
akhir = awal + 100
|
|
|
|
koleksi.upsert(
|
|
ids=ids[awal:akhir],
|
|
documents=documents[awal:akhir],
|
|
metadatas=metadatas[awal:akhir]
|
|
)
|
|
|
|
print(
|
|
f"{pdf.name}: {len(ids)} chunk "
|
|
f"untuk ukuran {ukuran}"
|
|
)
|
|
|
|
total_chunk += len(ids)
|
|
|
|
else:
|
|
print(
|
|
f"{pdf.name}: tidak ada teks yang terbaca."
|
|
)
|
|
|
|
print(f"Total chunk: {total_chunk}")
|
|
|
|
print("\nSemua indeks berhasil dibuat!")
|
|
print(f"Lokasi database: {FOLDER_DATABASE}")
|
|
|
|
|
|
# ==========================================
|
|
# PROGRAM UTAMA
|
|
# ==========================================
|
|
|
|
if __name__ == "__main__":
|
|
buat_indeks()
|