build_index.py

"""
build_index.py — Extract text from SQL PDFs, chunk it, embed it, and save a FAISS index.
Run this once locally before deploying to Hugging Face.
"""

import json
import os
import fitz  # PyMuPDF
import numpy as np
from sentence_transformers import SentenceTransformer

# ── Configuration ──────────────────────────────────────────────────────────────
DATA_DIR = os.path.join(os.path.dirname(__file__), "data")
INDEX_DIR = os.path.join(os.path.dirname(__file__), "faiss_index")
CHUNK_SIZE = 500       # approximate tokens (≈words for English)
CHUNK_OVERLAP = 50     # overlap between consecutive chunks
EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"


# ── PDF Extraction ─────────────────────────────────────────────────────────────
def extract_text_from_pdf(pdf_path: str) -> str:
    """Extract all text from a PDF using PyMuPDF."""
    doc = fitz.open(pdf_path)
    text = ""
    for page in doc:
        text += page.get_text()
    doc.close()
    return text


# ── Chunking ───────────────────────────────────────────────────────────────────
def chunk_text(text: str, source: str, chunk_size: int = CHUNK_SIZE, overlap: int = CHUNK_OVERLAP):
    """Split text into overlapping word-level chunks with metadata."""
    words = text.split()
    chunks = []
    start = 0
    while start < len(words):
        end = start + chunk_size
        chunk_words = words[start:end]
        chunk_text_str = " ".join(chunk_words)
        # Skip very short chunks (< 30 words)
        if len(chunk_words) >= 30:
            chunks.append({
                "text": chunk_text_str,
                "source": source,
                "chunk_id": len(chunks),
            })
        start += chunk_size - overlap
    return chunks


# ── Main ───────────────────────────────────────────────────────────────────────
def main():
    os.makedirs(INDEX_DIR, exist_ok=True)

    # 1. Extract & chunk all PDFs
    all_chunks = []
    pdf_files = [f for f in os.listdir(DATA_DIR) if f.lower().endswith(".pdf")]
    print(f"Found {len(pdf_files)} PDFs in {DATA_DIR}")

    for pdf_file in sorted(pdf_files):
        pdf_path = os.path.join(DATA_DIR, pdf_file)
        print(f"  Processing: {pdf_file} ...", end=" ", flush=True)
        text = extract_text_from_pdf(pdf_path)
        chunks = chunk_text(text, source=pdf_file)
        all_chunks.extend(chunks)
        print(f"{len(chunks)} chunks")

    print(f"\nTotal chunks: {len(all_chunks)}")

    # 2. Generate embeddings
    print(f"\nLoading embedding model: {EMBED_MODEL}")
    model = SentenceTransformer(EMBED_MODEL)

    texts = [c["text"] for c in all_chunks]
    print("Generating embeddings ...")
    embeddings = model.encode(texts, show_progress_bar=True, batch_size=64)
    embeddings = np.array(embeddings).astype("float32")
    print(f"Embeddings shape: {embeddings.shape}")

    # 3. Build FAISS index
    import faiss

    dimension = embeddings.shape[1]
    index = faiss.IndexFlatL2(dimension)
    index.add(embeddings)
    print(f"FAISS index size: {index.ntotal} vectors, dimension {dimension}")

    # 4. Save
    faiss.write_index(index, os.path.join(INDEX_DIR, "index.faiss"))
    with open(os.path.join(INDEX_DIR, "chunks.json"), "w", encoding="utf-8") as f:
        json.dump(all_chunks, f, ensure_ascii=False, indent=2)

    print(f"\n[OK] Index saved to {INDEX_DIR}/index.faiss")
    print(f"[OK] Chunks saved to {INDEX_DIR}/chunks.json")


if __name__ == "__main__":
    main()