File size: 4,295 Bytes
cd5f9d4 08c3203 cd5f9d4 08c3203 f1347ae 08c3203 f1347ae d68637e f1347ae 08c3203 cd5f9d4 08c3203 cd5f9d4 08c3203 cd5f9d4 08c3203 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | """Step 05 — potion broad retrieval: the recall pass with a static embedder, over the bm25_prefilter's output.
Embeds the pool's ``text`` ONLINE with the configured static embedder (``embedder_broad``, default
potion-base-32M — Model2Vec, a static token→vector lookup, ~13K docs/s, no transformer forward), scores each
candidate by cosine against the JD's ``ideal_summary`` (the subjective embedding query — HyDE was dropped),
and keeps the top ``top_n``. This is what lets the expensive granite embedder later touch only a few hundred
rows instead of the whole pool.
Deterministic: a static lookup with no RNG, plus a stable score+id sort — the surviving ids and
``potion_score`` are byte-reproducible across runs.
"""
from __future__ import annotations
import os
import time
from functools import lru_cache
from typing import Any
import numpy as np
import polars as pl
from common.logging import get_logger, step
from config import Settings, load_settings
from models.llm import MODEL_DIR # shared in-image model cache dir (HF_HOME redirect)
@lru_cache(maxsize=1)
def _load_potion(repo: str) -> Any:
"""Load (downloading on first use) the static Model2Vec embedder; cached for the process."""
os.environ.setdefault("HF_HOME", str(MODEL_DIR)) # resolve weights inside the repo, offline at rank time
from model2vec import StaticModel # imported lazily — keeps the import graph light
MODEL_DIR.mkdir(parents=True, exist_ok=True)
log = get_logger("retrieve")
log.info("potion.load", repo=repo)
start = time.perf_counter()
# force_download=True is model2vec's DEFAULT — it skips the local-cache check entirely and always calls
# snapshot_download, i.e. network on every load even when fully cached. False restores cache-first,
# download-only-if-missing (see model2vec.persistence._resolve_folder).
try:
model = StaticModel.from_pretrained(repo, force_download=False)
except Exception as exc:
log.error("potion.load_failed", repo=repo, seconds=round(time.perf_counter() - start, 3), error=str(exc))
raise
log.info("potion.load_complete", repo=repo, seconds=round(time.perf_counter() - start, 3))
return model
def _embed_unit(model: Any, texts: list[str]) -> np.ndarray:
"""Encode texts → L2-normalized float matrix (zero rows clamp to norm 1 so cosine is 0, never NaN)."""
vectors = np.asarray(model.encode(texts), dtype=np.float32)
norms = np.linalg.norm(vectors, axis=1, keepdims=True)
norms[norms == 0] = 1.0
return vectors / norms
def _narrow_by_cosine(candidates: pl.DataFrame, scores: np.ndarray, top_n: int) -> pl.DataFrame:
"""Attach ``potion_score`` and keep the top ``top_n`` (stable: score desc, then candidate_id)."""
keep_n = min(max(1, top_n), len(scores))
scored = candidates.with_columns(pl.Series("potion_score", scores.astype(float)))
return scored.sort(["potion_score", "candidate_id"], descending=[True, False]).head(keep_n)
def potion_broad_retrieval(
candidates: pl.DataFrame, query_text: str, settings: Settings | None = None
) -> pl.DataFrame:
"""Narrow the pool to the top-``top_n`` by potion cosine vs ``query_text`` (the JD ideal_summary).
Adds a ``potion_score`` column. Disabled / empty query / empty pool ⇒ pass the full pool through with
``potion_score = 0.0`` (downstream stages then no-op on the flat column).
"""
active = settings or load_settings()
broad = active.broad_retrieval
with step("rank.potion_broad", rows_in=candidates.height) as metrics:
if not broad.enabled or not query_text.strip() or candidates.height == 0:
metrics["rows_out"] = candidates.height
metrics["skipped"] = True
return candidates.with_columns(pl.lit(0.0).alias("potion_score"))
model = _load_potion(active.models.embedder_broad)
pool_vectors = _embed_unit(model, candidates.get_column("text").to_list())
query_vector = _embed_unit(model, [query_text])[0]
scores = pool_vectors @ query_vector # cosine (both L2-normalized)
narrowed = _narrow_by_cosine(candidates, scores, broad.top_n)
metrics["rows_out"] = narrowed.height
metrics["score_max"] = round(float(scores.max()), 4)
return narrowed
|