scripts/build_index.py

"""Embed every chunk in data/chunks.jsonl and persist a FAISS index.

Usage:
    uv run python scripts/build_index.py
"""
from __future__ import annotations

import json
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from rich.console import Console

from researchpath.embeddings import Embedder
from researchpath.index import build_index

console = Console()
ROOT = Path(__file__).resolve().parents[1]
CHUNKS_PATH = ROOT / "data" / "chunks.jsonl"
INDEX_PATH = ROOT / "data" / "index.faiss"


def main() -> int:
    if not CHUNKS_PATH.exists():
        console.print(f"[red]Missing {CHUNKS_PATH}. Run scripts/parse_corpus.py first.[/red]")
        return 1

    chunks: list[dict] = []
    with open(CHUNKS_PATH, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                chunks.append(json.loads(line))

    console.print(f"[bold]Embedding {len(chunks)} chunks with BAAI/bge-small-en-v1.5[/bold]")
    console.print("[dim](first run downloads ~133MB model from HF Hub; subsequent runs are cached)[/dim]\n")

    t0 = time.time()
    embedder = Embedder()
    build_index(chunks, embedder, INDEX_PATH)
    dt = time.time() - t0

    console.print(f"\n[green]Wrote {INDEX_PATH.relative_to(ROOT)} in {dt:.1f}s[/green]")
    console.print(f"[green]Wrote {INDEX_PATH.with_suffix('.chunks.json').relative_to(ROOT)}[/green]")
    return 0


if __name__ == "__main__":
    sys.exit(main())