韓国語PDF検索に埋め込みとリランカーを追加する

python3 -m venv .venv
source .venv/bin/activate
python -m pip install -U sentence-transformers pypdf

from pathlib import Path
from pypdf import PdfReader
from sentence_transformers import SentenceTransformer, CrossEncoder
from sentence_transformers.util import cos_sim
from time import perf_counter

pdf_path = Path("sample.pdf")
reader = PdfReader(pdf_path)
passages = []
for page_number, page in enumerate(reader.pages, start=1):
    text = " ".join((page.extract_text() or "").split())
    for start_char in range(0, len(text), 700):
        chunk = text[start_char:start_char + 700].strip()
        if chunk:
            passages.append({"chunk_id": f"p{page_number}-{start_char}", "text": chunk, "source": pdf_path.name, "page": page_number})
if not passages:
    raise ValueError("No selectable PDF text found; use a text PDF or run OCR first.")

query = "환불은 며칠 안에 신청해야 하나요?"
embedder = SentenceTransformer("BAAI/bge-m3")
reranker = CrossEncoder("BAAI/bge-reranker-v2-m3")
# One-time indexing is outside both per-query timers.
doc_vectors = embedder.encode([p["text"] for p in passages], normalize_embeddings=True)
start = perf_counter()
query_vector = embedder.encode([query], normalize_embeddings=True)
candidate_scores = cos_sim(query_vector, doc_vectors)[0].tolist()
candidates = sorted(zip(passages, candidate_scores), key=lambda x: x[1], reverse=True)[:3]
baseline_ms = (perf_counter() - start) * 1000
print("baseline ranks:")
for rank, (p, score) in enumerate(candidates, 1):
    print(rank, p["source"], p["page"], p["chunk_id"], round(float(score), 4))

start = perf_counter()
r_scores = reranker.predict([(query, p["text"]) for p, _ in candidates])
rerank_ms = (perf_counter() - start) * 1000
reranked = sorted(zip([p for p, _ in candidates], r_scores), key=lambda x: float(x[1]), reverse=True)
print("reranked ranks:")
for rank, (p, score) in enumerate(reranked, 1):
    print(rank, p["source"], p["page"], p["chunk_id"], round(float(score), 4))
print("baseline_query_ms", round(baseline_ms, 2))
print("reranker_only_ms", round(rerank_ms, 2))
print("query_total_ms", round(baseline_ms + rerank_ms, 2))

python rag_compare.py