scripts/build_fast_indexes.py

from __future__ import annotations

import argparse
import json
import logging
import sys
import time
from pathlib import Path

import faiss
import numpy as np
from rank_bm25 import BM25Okapi
from sklearn.feature_extraction.text import HashingVectorizer
from tqdm import tqdm

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from src.core.config import DATA_DIR
from src.core.constants import CANDIDATES_PATH
from src.ingestion.normalizer import normalize_redrob
from src.ingestion.parser import ProfileParser

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def _build_document_text(profile) -> str:
    parts: list[str] = []
    if profile.raw_text:
        parts.append(profile.raw_text)
    parts.extend(s.name for s in profile.skills)
    for exp in profile.experience:
        parts.append(exp.title)
        parts.append(exp.company)
        parts.append(exp.description)
    for edu in profile.education:
        parts.append(edu.institution)
        if edu.field:
            parts.append(edu.field)
    if profile.professional and profile.professional.current_title:
        parts.append(profile.professional.current_title)
    if profile.professional and profile.professional.current_company:
        parts.append(profile.professional.current_company)
    return " ".join(p for p in parts if p)


def _tokenize(text: str) -> list[str]:
    return text.lower().split()


def _save_offset_index(profiles_path: Path, profile_ids: set[str], output_path: Path) -> None:
    offsets: dict[str, int] = {}
    with open(profiles_path, encoding="utf-8") as f:
        while True:
            offset = f.tell()
            line = f.readline()
            if not line:
                break
            line = line.strip()
            if not line:
                continue
            try:
                raw = json.loads(line)
            except json.JSONDecodeError:
                continue
            cand_id = raw.get("profile_id") or raw.get("candidate_id") or raw.get("id")
            profile_nested = raw.get("profile", {})
            if isinstance(profile_nested, dict) and not cand_id:
                cand_id = (
                    profile_nested.get("profile_id")
                    or profile_nested.get("candidate_id")
                    or profile_nested.get("id")
                )
            if cand_id and str(cand_id) in profile_ids:
                offsets[str(cand_id)] = offset

    output_path.parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, "w") as f:
        json.dump(offsets, f)
    logger.info("Offset index saved: %d entries", len(offsets))


def build_fast_indexes(profiles_path: Path = CANDIDATES_PATH, batch_size: int = 2000) -> None:
    start = time.perf_counter()
    index_dir = DATA_DIR / "indexes"
    index_dir.mkdir(parents=True, exist_ok=True)

    faiss_path = index_dir / "faiss_index.bin"
    id_map_path = index_dir / "faiss_id_map.json"
    bm25_path = index_dir / "bm25_index.pkl"
    offset_path = index_dir / "offset_index.json"
    meta_path = index_dir / "index_meta.json"

    parser = ProfileParser()
    vectorizer = HashingVectorizer(
        n_features=384,
        alternate_sign=False,
        norm="l2",
        lowercase=True,
        token_pattern=r"(?u)\b\w+\b",
    )
    index = faiss.IndexHNSWFlat(384, 32, faiss.METRIC_INNER_PRODUCT)
    profile_ids: list[str] = []
    corpus_tokenized: list[list[str]] = []

    batch_texts: list[str] = []
    loaded = 0
    skipped = 0

    def flush_batch() -> None:
        if not batch_texts:
            return
        sparse = vectorizer.transform(batch_texts)
        dense = sparse.astype(np.float32).toarray()
        index.add(dense)
        batch_texts.clear()

    logger.info("Building fast indexes from %s", profiles_path)
    for raw in tqdm(parser.parse_jsonl_file(profiles_path), desc="Profiles", unit="profile"):
        try:
            profile = normalize_redrob(raw)
            text = _build_document_text(profile)
            profile_ids.append(profile.profile_id)
            corpus_tokenized.append(_tokenize(text))
            batch_texts.append(text)
            loaded += 1
            if len(batch_texts) >= batch_size:
                flush_batch()
        except Exception:
            skipped += 1

    flush_batch()

    logger.info("Loaded %d profiles (%d skipped)", loaded, skipped)
    faiss.write_index(index, str(faiss_path))
    with open(id_map_path, "w") as f:
        json.dump(profile_ids, f)
    with open(meta_path, "w") as f:
        json.dump({"embedding": "hashing", "dimension": 384}, f)
    logger.info("FAISS index saved: %d vectors", index.ntotal)

    import pickle
    with open(bm25_path, "wb") as f:
        pickle.dump({"corpus_tokenized": corpus_tokenized, "id_map": profile_ids}, f)
    logger.info("BM25 index saved: %d documents", len(profile_ids))

    _save_offset_index(profiles_path, set(profile_ids), offset_path)
    logger.info("All fast indexes built in %.1fs", time.perf_counter() - start)


def main() -> None:
    parser = argparse.ArgumentParser(description="Build fast full-dataset indexes")
    parser.add_argument("--profiles", type=Path, default=CANDIDATES_PATH)
    parser.add_argument("--batch-size", type=int, default=2000)
    args = parser.parse_args()
    build_fast_indexes(args.profiles, args.batch_size)


if __name__ == "__main__":
    main()