Mengindeks Set Dokumen

# PDFs
from pypdf import PdfReader
text = ''
for page in PdfReader('doc.pdf').pages:
    text += page.extract_text()

# HTML
from bs4 import BeautifulSoup
text = BeautifulSoup(html, 'html.parser').get_text()

# Markdown — just read the file
text = open('doc.md').read()

from langchain.text_splitter import RecursiveCharacterTextSplitter
splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100)
chunks = splitter.split_text(text)

from openai import OpenAI
client = OpenAI()

def embed_batch(texts, batch_size=100):
    vectors = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]
        resp = client.embeddings.create(model='text-embedding-3-small', input=batch)
        vectors.extend(d.embedding for d in resp.data)
    return vectors

import chromadb
client = chromadb.Client()
collection = client.create_collection('docs')

collection.add(
    ids=[f'doc-{i}' for i in range(len(chunks))],
    embeddings=vectors,
    documents=chunks,
    metadatas=[{'source': 'doc.pdf', 'page': i} for i in range(len(chunks))]
)

import hashlib

def chunk_id(source, text):
    h = hashlib.sha256(text.encode()).hexdigest()[:16]
    return f'{source}:{h}'

print(chunk_id('handbook.pdf', 'Employees get 20 vacation days per year.'))
print(chunk_id('handbook.pdf', 'Employees get 20 vacation days per year.'))

metadata = {
    'source': '/docs/handbook.pdf',
    'page': 42,
    'department': 'engineering',
    'updated_at': '2024-08-12',
    'access_level': 'internal'
}
for k, v in metadata.items():
    print(f"{k}: {v}")

for i, batch in enumerate(batches):
    embed_and_store(batch)
    if i % 10 == 0:
        save_checkpoint(i)
        print(f'Processed {i * 100} chunks')

from asyncio import Semaphore
sem = Semaphore(10)  # 10 concurrent embed calls max

async def safe_embed(batch):
    async with sem:
        return await client.embeddings.create(...)

collection = client.create_collection(f'docs-v2-{date.today()}')
# old collection stays live until new one is validated