Spaces:
Runtime error
Runtime error
Download scripts/preprocess.py from ashokbugude/redrob-ranker: direct link, hf CLI and curl.
- Browser
- Download file 14.9 kB
-
https://huggingface.co/spaces/ashokbugude/redrob-ranker/resolve/main/scripts/preprocess.py
- Command line
-
hf download hf://spaces/ashokbugude/redrob-ranker/scripts/preprocess.py
-
curl -L -o preprocess.py https://huggingface.co/spaces/ashokbugude/redrob-ranker/resolve/main/scripts/preprocess.py
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()) | |