from __future__ import annotations import argparse import logging import sys import time from pathlib import Path 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, SAMPLE_PATH from src.core.models import Profile from src.ingestion.normalizer import normalize_redrob from src.ingestion.parser import ProfileParser from src.ingestion.quality_scorer import compute_data_quality_score from src.language.multilingual import MultilingualEmbedder from src.search.bm25_search import BM25Search from src.search.vector_search import VectorSearch logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) def build_indexes( profiles_path: Path = CANDIDATES_PATH, sample_count: int = 0, force: bool = False, ) -> None: start = time.perf_counter() if not profiles_path.exists(): logger.warning(f"Profiles file not found: {profiles_path}") logger.info(f"Falling back to sample data: {SAMPLE_PATH}") profiles_path = SAMPLE_PATH if not profiles_path.exists(): logger.error(f"No data found at {profiles_path} or {CANDIDATES_PATH}") return index_dir = DATA_DIR / "indexes" faiss_path = index_dir / "faiss_index.bin" bm25_path = index_dir / "bm25_index.pkl" if faiss_path.exists() and bm25_path.exists() and not force: logger.info("Indexes already exist. Use --force to rebuild.") return logger.info(f"Loading profiles from {profiles_path}") parser = ProfileParser() profiles: list[Profile] = [] loaded = 0 skipped = 0 if profiles_path.suffix == ".jsonl": log_interval = max(1, sample_count // 10) if sample_count > 0 else 10000 for raw in parser.parse_jsonl_file(profiles_path): try: normalized = normalize_redrob(raw) qs = compute_data_quality_score(normalized) if qs < 0.3: skipped += 1 continue profiles.append(normalized) loaded += 1 if loaded % log_interval == 0 and sample_count == 0: logger.info(f" Loaded {loaded} profiles ({skipped} skipped so far)...") except Exception: skipped += 1 continue if sample_count > 0 and loaded >= sample_count: break else: data = parser.parse_json_file(profiles_path) for item in data: try: normalized = normalize_redrob(item) profiles.append(normalized) loaded += 1 except Exception: skipped += 1 logger.info(f"Loaded {len(profiles)} profiles ({skipped} skipped)") if not profiles: logger.error("No valid profiles to index") return raw_texts = [p.raw_text for p in profiles] profile_ids = [p.profile_id for p in profiles] document_texts = [_build_document_text(p) for p in profiles] logger.info("Generating embeddings (this may take a while)...") embedder = MultilingualEmbedder() batch_size = 500 all_embeddings = [] num_batches = (len(raw_texts) + batch_size - 1) // batch_size with tqdm(total=num_batches, desc="Embedding", unit="batch") as pbar: for i in range(0, len(raw_texts), batch_size): batch = raw_texts[i : i + batch_size] batch_emb = embedder.embed_batch(batch) all_embeddings.append(batch_emb) pbar.update(1) import numpy as np embeddings = np.vstack(all_embeddings) if len(all_embeddings) > 1 else all_embeddings[0] logger.info(f"Generated {len(embeddings)} embeddings (dim={embeddings.shape[1]})") vector_search = VectorSearch(dimension=384) vector_search.build_index(embeddings, profile_ids) vector_search.save() logger.info(f"FAISS index saved: {vector_search.size} vectors") bm25_search = BM25Search() bm25_search.build_index(document_texts, profile_ids) bm25_search.save() logger.info(f"BM25 index saved: {bm25_search.size} documents") offset_path = index_dir / "offset_index.json" _save_offset_index(profiles_path, profile_ids, offset_path) elapsed = time.perf_counter() - start logger.info(f"All indexes built successfully in {elapsed:.1f}s") def _build_document_text(profile: 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(parts) def _save_offset_index( profiles_path: Path, profile_ids: list[str], output_path: Path, ) -> None: import json if profiles_path.suffix != ".jsonl": logger.info("Skipping offset index (not a JSONL file)") return pid_set = set(profile_ids) 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) 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 pid_set: offsets[str(cand_id)] = offset except json.JSONDecodeError: continue output_path.parent.mkdir(parents=True, exist_ok=True) with open(output_path, "w") as f: json.dump(offsets, f) logger.info(f"Offset index saved: {len(offsets)} entries") def main(): parser = argparse.ArgumentParser(description="Build FAISS + BM25 indexes from profiles") parser.add_argument("--profiles", type=str, default=None, help="Path to candidate profiles file (JSON or JSONL)") parser.add_argument("--sample", type=int, default=0, help="Process only N profiles (for quick testing)") parser.add_argument("--force", action="store_true", help="Rebuild indexes even if they exist") args = parser.parse_args() profiles_p = Path(args.profiles) if args.profiles else CANDIDATES_PATH build_indexes(profiles_path=profiles_p, sample_count=args.sample, force=args.force) if __name__ == "__main__": main()