38 lines
1.3 KiB
Python
38 lines
1.3 KiB
Python
from pathlib import Path
|
|
from pypdf import PdfReader
|
|
import chromadb
|
|
from chromadb.utils import embedding_functions
|
|
|
|
def baca_pdf(path):
|
|
r = PdfReader(path)
|
|
return "\n".join((p.extract_text() or "") for p in r.pages)
|
|
|
|
def potong(teks, ukuran=500, tumpang=75):
|
|
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()]
|
|
|
|
print("Menyiapkan model embedding (unduh jika belum ada)...")
|
|
# Model multilingual agar bisa bahasa Indonesia
|
|
emb = embedding_functions.SentenceTransformerEmbeddingFunction(
|
|
model_name="paraphrase-multilingual-MiniLM-L12-v2")
|
|
|
|
client = chromadb.PersistentClient(path="./chroma_db")
|
|
koleksi = client.get_or_create_collection(
|
|
name="pedoman", embedding_function=emb,
|
|
metadata={"hnsw:space": "cosine"})
|
|
|
|
# Proses tiap PDF
|
|
for pdf in Path("dokumen").glob("*.pdf"):
|
|
teks = baca_pdf(pdf)
|
|
chunks = potong(teks, ukuran=500, tumpang=75) # Pakai ukuran 500
|
|
koleksi.upsert(
|
|
ids=[f"{pdf.stem}-{i}" for i in range(len(chunks))],
|
|
documents=chunks,
|
|
metadatas=[{"sumber": pdf.name, "chunk": i} for i in range(len(chunks))],
|
|
)
|
|
print(f"{pdf.name} → {len(chunks)} chunk disimpan.")
|
|
|
|
print("\nSelesai! Indeks tersimpan di folder 'chroma_db'.") |