redrob-ranker / scripts /preprocess.py
ashokbugude
Cursor
Full HF Space: bundled candidates.jsonl, live rank.py, optional upload
b91193c
Raw History Blame Contribute Delete
14.9 kB
#!/usr/bin/env python3
"""Offline preprocessing orchestrator β€” Phase 2."""
from __future__ import annotations
import argparse
import json
import logging
import sys
import time
from collections.abc import Callable
from pathlib import Path
PROJECT_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(PROJECT_ROOT))
from app.artifacts_manifest import is_locked, write_manifest # noqa: E402
from app.bm25_index import build_bm25_index, save_bm25 # noqa: E402
from app.candidates import count_candidates, load_candidates, load_candidates_list # noqa: E402
from app.config import get_settings # noqa: E402
from app.embeddings import encode_candidates, save_embeddings # noqa: E402
from app.feature_store import ( # noqa: E402
build_features_frame,
load_silver_tiers,
validate_features_frame,
write_features_parquet,
)
from app.features import FEATURE_NAMES # noqa: E402
from app.jd_requirements import load_or_build_jd_requirements # noqa: E402
from app.gemini_client import gemini_jd_enabled, gemini_labels_enabled # noqa: E402
from app.logging_setup import configure_logging # noqa: E402
from app.progress import format_duration # noqa: E402
logger = logging.getLogger(__name__)
STEPS = ("jd", "features", "embeddings", "bm25", "labels", "career_recall", "all")
def _candidate_total(candidates_path: Path, limit: int | None) -> int:
if limit is not None:
return limit
return count_candidates(candidates_path)
def _run_career_recall(settings, candidates_path: Path, artifacts_dir: Path) -> None:
from scripts.build_career_recall_scores import build_scores, TOP_K, MIN_SCORE
import json
from app.artifact_names import CAREER_SCORES
out_path = artifacts_dir / CAREER_SCORES
scores = build_scores(candidates_path, settings=settings)
payload = {"scores": scores, "top_k": TOP_K, "min_score": MIN_SCORE}
out_path.write_text(json.dumps(payload, indent=2), encoding="utf-8")
logger.info("Wrote %d career recall scores to %s", len(scores), out_path)
def _features_stale(settings) -> bool:
"""True when JD was rebuilt after the feature parquet (skill/title features depend on JD)."""
feat_path = settings.artifact_path("candidate_features.parquet")
jd_path = settings.jd_requirements_path
if not feat_path.exists() or not jd_path.exists():
return False
return jd_path.stat().st_mtime_ns > feat_path.stat().st_mtime_ns
def run_jd(settings, artifacts_dir: Path, *, force: bool = False, force_llm: bool = False) -> None:
jd_path = settings.jd_requirements_path
if jd_path.exists() and not force:
if gemini_jd_enabled(settings):
try:
cached = json.loads(jd_path.read_text(encoding="utf-8"))
if cached.get("source") != "gemini_pro":
logger.info(
"Cached JD source=%s but Gemini JD enabled β€” rebuilding with Gemini Pro",
cached.get("source", "?"),
)
force = True
else:
logger.info("Using cached gemini_pro JD requirements at %s", jd_path)
return
except (OSError, json.JSONDecodeError):
force = True
else:
logger.info("Using cached JD requirements at %s", jd_path)
return
if not gemini_jd_enabled(settings):
logger.info("Gemini JD parse disabled β€” building heuristic JD requirements")
load_or_build_jd_requirements(
settings.job_description_path,
settings.jd_requirements_path,
settings=settings,
force=force,
)
def run_features(settings, candidates_path: Path, *, limit: int | None = None, force: bool = False) -> None:
out_path = settings.artifact_path("candidate_features.parquet")
if out_path.exists() and not force:
if _features_stale(settings):
logger.info("JD requirements newer than features β€” rebuilding feature parquet")
else:
logger.info("Features exist at %s β€” skipping (use --force)", out_path)
return
total = _candidate_total(candidates_path, limit)
logger.info("Feature extraction: %d candidates", total)
jd = load_or_build_jd_requirements(
settings.job_description_path,
settings.jd_requirements_path,
settings=settings,
)
silver = load_silver_tiers(settings.artifact_path("labels_silver.parquet"))
gemini_path = settings.artifact_path("gemini_tiers.parquet")
gemini: dict[str, int] = {}
if gemini_path.exists():
import polars as pl
gdf = pl.read_parquet(gemini_path)
gemini = dict(zip(gdf["candidate_id"].to_list(), gdf["gemini_tier"].to_list(), strict=False))
candidates = load_candidates(candidates_path, limit=limit)
frame = build_features_frame(
candidates,
jd,
silver_tiers=silver,
gemini_tiers=gemini,
total_candidates=total,
)
validate_features_frame(frame)
write_features_parquet(frame, out_path)
logger.info("Feature columns: %d features + metadata", len(FEATURE_NAMES))
def run_embeddings(
settings,
candidates_path: Path,
*,
limit: int | None = None,
force: bool = False,
device: str | None = None,
batch_size: int | None = None,
) -> None:
npy_path = settings.artifact_path("bge_embeddings.npy")
index_path = settings.artifact_path("faiss.index")
if npy_path.exists() and index_path.exists() and not force:
logger.info("Embeddings exist β€” skipping (use --force)")
return
from app.embeddings import default_batch_size, resolve_embedding_device
candidates = load_candidates_list(candidates_path, limit=limit)
resolved = resolve_embedding_device(device)
effective_batch = batch_size or default_batch_size(resolved)
total = len(candidates)
batches = (total + effective_batch - 1) // effective_batch if total else 0
logger.info(
"Embeddings: encoding %d candidates on %s (batch_size=%d, ~%d batches, tqdm bar below)",
total,
resolved,
effective_batch,
batches,
)
t0 = time.perf_counter()
matrix, ids = encode_candidates(
candidates,
show_progress=total > 100,
device=resolved,
batch_size=batch_size,
)
logger.info("Embeddings: encode finished in %s", format_duration(time.perf_counter() - t0))
t1 = time.perf_counter()
save_embeddings(matrix, ids, npy_path=npy_path, index_path=index_path)
logger.info("Embeddings: saved FAISS + npy in %s", format_duration(time.perf_counter() - t1))
def run_bm25(settings, candidates_path: Path, *, limit: int | None = None, force: bool = False) -> None:
bm25_path = settings.artifact_path("bm25.pkl")
if bm25_path.exists() and not force:
logger.info("BM25 index exists β€” skipping (use --force)")
return
total = _candidate_total(candidates_path, limit)
logger.info("BM25: building index for %d candidates", total)
candidates = load_candidates_list(candidates_path, limit=limit)
artifacts = build_bm25_index(candidates)
save_bm25(artifacts, bm25_path)
def run_labels(
settings,
candidates_path: Path,
artifacts_dir: Path,
*,
limit: int | None = None,
force: bool = False,
force_llm: bool = False,
) -> None:
from scripts.label_archetypes import build_gemini_tiers
out_path = settings.artifact_path("gemini_tiers.parquet")
if out_path.exists() and not force and not force_llm:
logger.info("Gemini tiers exist at %s β€” skipping (use --force-llm to re-label)", out_path)
return
total = _candidate_total(candidates_path, limit)
if gemini_labels_enabled(settings):
logger.info("Gemini archetype labels: %d candidates (API enabled)", total)
else:
logger.info("Silver-tier labels only: %d candidates (Gemini labels disabled)", total)
jd = load_or_build_jd_requirements(
settings.job_description_path,
settings.jd_requirements_path,
settings=settings,
)
build_gemini_tiers(
candidates_path,
out_path,
settings=settings,
jd=jd,
limit=limit,
force=force or force_llm,
skip_api=not gemini_labels_enabled(settings) or (is_locked(artifacts_dir) and not force_llm),
)
def _run_named_step(
name: str,
fn: Callable[[], None],
*,
index: int | None = None,
total_steps: int | None = None,
) -> None:
if index is not None and total_steps is not None:
logger.info("=== Pipeline step %d/%d: %s ===", index, total_steps, name)
else:
logger.info("=== Step: %s ===", name)
t0 = time.perf_counter()
fn()
logger.info("=== %s complete (elapsed %s) ===", name, format_duration(time.perf_counter() - t0))
def main() -> int:
parser = argparse.ArgumentParser(description="Run offline preprocessing pipeline (Phase 2).")
parser.add_argument(
"--step",
choices=STEPS,
default="all",
help="Pipeline step to run (default: all).",
)
parser.add_argument("--candidates", type=Path, default=None)
parser.add_argument("--limit", type=int, default=None, help="Limit candidates for quick runs.")
parser.add_argument("--force", action="store_true", help="Rebuild even if artifacts exist.")
parser.add_argument(
"--skip-llm",
action="store_true",
help="Disable all Gemini calls (same as --no-gemini-jd --no-gemini-labels).",
)
parser.add_argument(
"--gemini-jd",
action=argparse.BooleanOptionalAction,
default=None,
help="Enable/disable Gemini Pro JD parse (default: from REDROB_GEMINI_JD_PARSE).",
)
parser.add_argument(
"--gemini-labels",
action=argparse.BooleanOptionalAction,
default=None,
help="Enable/disable Gemini Flash 100K labels (default: from REDROB_GEMINI_LABELS).",
)
parser.add_argument(
"--force-llm",
action="store_true",
help="Re-run Gemini archetype labels even when gemini_tiers.parquet exists.",
)
parser.add_argument(
"--device",
choices=("cuda", "cpu"),
default=None,
help="Embedding device override (default: auto-detect CUDA).",
)
parser.add_argument(
"--batch-size",
type=int,
default=None,
help="Embedding batch size override (default: 256 on CUDA, 64 on CPU).",
)
args = parser.parse_args()
settings = get_settings()
settings_updates: dict = {}
if args.skip_llm:
settings_updates["skip_gemini"] = True
if args.gemini_jd is not None:
settings_updates["gemini_jd_parse"] = args.gemini_jd
if args.gemini_labels is not None:
settings_updates["gemini_labels"] = args.gemini_labels
if settings_updates:
settings = settings.model_copy(update=settings_updates)
configure_logging(settings.log_level)
settings.ensure_artifacts_dir()
artifacts_dir = settings.artifacts_dir.resolve()
candidates_path = (args.candidates or settings.candidates_path).resolve()
step = args.step
pipeline_t0 = time.perf_counter()
if step == "all":
total_candidates = _candidate_total(candidates_path, args.limit)
logger.info(
"Preprocess pipeline: %d candidates | gemini_jd=%s | gemini_labels=%s | force=%s | force_llm=%s",
total_candidates,
gemini_jd_enabled(settings),
gemini_labels_enabled(settings),
args.force,
args.force_llm,
)
planned: list[tuple[str, Callable[[], None]]] = [
("JD parse", lambda: run_jd(settings, artifacts_dir, force=args.force, force_llm=args.force_llm)),
(
"Gemini archetype labels",
lambda: run_labels(
settings,
candidates_path,
artifacts_dir,
limit=args.limit,
force=args.force,
force_llm=args.force_llm,
),
),
(
"Feature extraction",
lambda: run_features(settings, candidates_path, limit=args.limit, force=args.force),
),
(
"BGE embeddings + FAISS",
lambda: run_embeddings(
settings,
candidates_path,
limit=args.limit,
force=args.force,
device=args.device,
batch_size=args.batch_size,
),
),
("BM25 index", lambda: run_bm25(settings, candidates_path, limit=args.limit, force=args.force)),
]
for idx, (name, fn) in enumerate(planned, start=1):
_run_named_step(name, fn, index=idx, total_steps=len(planned))
if not args.force_llm:
write_manifest(artifacts_dir, locked=is_locked(artifacts_dir))
logger.info(
"Preprocess pipeline complete β€” total elapsed %s",
format_duration(time.perf_counter() - pipeline_t0),
)
return 0
if step == "jd":
_run_named_step("JD parse", lambda: run_jd(settings, artifacts_dir, force=args.force, force_llm=args.force_llm))
elif step == "labels":
_run_named_step(
"Gemini archetype labels",
lambda: run_labels(
settings,
candidates_path,
artifacts_dir,
limit=args.limit,
force=args.force,
force_llm=args.force_llm,
),
)
elif step == "features":
_run_named_step(
"Feature extraction",
lambda: run_features(settings, candidates_path, limit=args.limit, force=args.force),
)
elif step == "embeddings":
_run_named_step(
"BGE embeddings + FAISS",
lambda: run_embeddings(
settings,
candidates_path,
limit=args.limit,
force=args.force,
device=args.device,
batch_size=args.batch_size,
),
)
elif step == "bm25":
_run_named_step(
"BM25 index",
lambda: run_bm25(settings, candidates_path, limit=args.limit, force=args.force),
)
elif step == "career_recall":
_run_named_step(
"Career recall scores",
lambda: _run_career_recall(settings, candidates_path, artifacts_dir),
)
logger.info("Preprocessing step '%s' complete (elapsed %s).", step, format_duration(time.perf_counter() - pipeline_t0))
return 0
if __name__ == "__main__":
raise SystemExit(main())