46 lines
1.9 KiB
Python
46 lines
1.9 KiB
Python
"""Lab 2 - Bangun indeks ChromaDB dari semua PDF di folder dokumen/.
|
|
|
|
Pakai:
|
|
python indeks.py # ukuran 800, overlap 100 (koleksi: pedoman_800)
|
|
python indeks.py 500 1000 # beberapa ukuran sekaligus (overlap 15%)
|
|
Tiap ukuran menjadi koleksi sendiri (pedoman_500, pedoman_1000, ...) untuk Lab 3.
|
|
Aman dijalankan ulang (upsert).
|
|
"""
|
|
import sys
|
|
from umum import (FOLDER_DOKUMEN, baca_pdf_per_halaman, potong, nama_koleksi,
|
|
ambil_embedding, buka_klien, buat_koleksi)
|
|
|
|
if len(sys.argv) > 1:
|
|
daftar = [(int(a), int(int(a) * 0.15)) for a in sys.argv[1:]]
|
|
else:
|
|
daftar = [(800, 100)]
|
|
|
|
pdfs = sorted(FOLDER_DOKUMEN.glob("*.pdf"))
|
|
if not pdfs:
|
|
sys.exit("Tidak ada PDF di folder dokumen/. Letakkan 2-3 PDF berbasis teks di sana.")
|
|
|
|
print("Memuat model embedding (pertama kali akan mengunduh ratusan MB)...")
|
|
emb = ambil_embedding()
|
|
klien = buka_klien()
|
|
|
|
for ukuran, tumpang in daftar:
|
|
koleksi = buat_koleksi(klien, nama_koleksi(ukuran), emb)
|
|
print(f"\n== Koleksi {koleksi.name} (ukuran={ukuran}, overlap={tumpang}) ==")
|
|
for pdf in pdfs:
|
|
halaman = baca_pdf_per_halaman(pdf)
|
|
if not halaman:
|
|
print(f"{pdf.name}: TEKS KOSONG (PDF scan?) - dilewati")
|
|
continue
|
|
ids, docs, metas = [], [], []
|
|
n = 0
|
|
for no_hal, teks in halaman: # potong per halaman agar metadata halaman akurat
|
|
for c in potong(teks, ukuran, tumpang):
|
|
ids.append(f"{pdf.stem}-{ukuran}-{n}")
|
|
docs.append(c)
|
|
metas.append({"sumber": pdf.name, "halaman": no_hal, "chunk": n})
|
|
n += 1
|
|
for i in range(0, len(ids), 100): # kirim bertahap
|
|
koleksi.upsert(ids=ids[i:i+100], documents=docs[i:i+100], metadatas=metas[i:i+100])
|
|
print(f"{pdf.name} -> {len(ids)} chunk")
|
|
print(f"Total di koleksi: {koleksi.count()} chunk")
|