Spaces:
Runtime error
Runtime error
| from __future__ import annotations | |
| from typing import Dict, List, Tuple | |
| import pandas as pd | |
| from schemas.candidates import Candidate, CandidateSet | |
| def retrieve_candidates(plan, bm25, vector, topn: int = 200, catalog_df: pd.DataFrame | None = None) -> CandidateSet: | |
| """Retrieve union of BM25 + vector using plan queries, then apply weighted RRF fusion.""" | |
| bm25_res = bm25.recommend(plan.bm25_query, k=topn, return_scores=True) | |
| vec_res = vector.recommend(plan.vec_query, k=topn, return_scores=True) | |
| def to_id_scores(res): | |
| out_ids = [] | |
| out_scores = {} | |
| for r in res: | |
| if isinstance(r, dict) and "assessment_id" in r: | |
| aid = r["assessment_id"] | |
| out_ids.append(aid) | |
| if "score" in r: | |
| out_scores[aid] = r["score"] | |
| elif isinstance(r, (tuple, list)) and len(r) >= 2: | |
| aid, sc = r[0], r[1] | |
| out_ids.append(aid) | |
| try: | |
| out_scores[aid] = float(sc) | |
| except Exception: | |
| pass | |
| elif isinstance(r, str): | |
| out_ids.append(r) | |
| return out_ids, out_scores | |
| bm25_ids, bm25_scores = to_id_scores(bm25_res) | |
| vec_ids, vec_scores = to_id_scores(vec_res) | |
| bm25_pos = {aid: i + 1 for i, aid in enumerate(bm25_ids)} | |
| vec_pos = {aid: i + 1 for i, aid in enumerate(vec_ids)} | |
| # union preserving ids | |
| seen = set() | |
| union_ids: List[str] = [] | |
| for aid in bm25_ids: | |
| if aid not in seen: | |
| union_ids.append(aid) | |
| seen.add(aid) | |
| for aid in vec_ids: | |
| if aid not in seen: | |
| union_ids.append(aid) | |
| seen.add(aid) | |
| # Query-adaptive weights + RRF fusion | |
| def _choose_fusion_weights(plan, raw_query: str) -> Tuple[float, float]: | |
| q = (plan.rerank_query or raw_query or "").strip() | |
| n_words = len(q.split()) | |
| n_skills = len(plan.must_have_skills or []) | |
| n_soft = len(plan.soft_skills or []) | |
| w_b, w_v = 0.5, 0.5 | |
| if n_skills >= 2: | |
| w_b += 0.2 | |
| w_v -= 0.2 | |
| if n_words >= 18: | |
| w_v += 0.2 | |
| w_b -= 0.2 | |
| if n_soft >= 2: | |
| w_v += 0.1 | |
| w_b -= 0.1 | |
| w_b = max(0.1, min(0.9, w_b)) | |
| w_v = max(0.1, min(0.9, w_v)) | |
| s = w_b + w_v | |
| return w_b / s, w_v / s | |
| w_bm25, w_vec = _choose_fusion_weights(plan, raw_query=plan.bm25_query) | |
| k_rrf = 60.0 | |
| candidates: List[Candidate] = [] | |
| scored: List[Tuple[float, Candidate]] = [] | |
| for aid in union_ids: | |
| rb = bm25_pos.get(aid) | |
| rv = vec_pos.get(aid) | |
| rrf_b = w_bm25 / (k_rrf + rb) if rb is not None else 0.0 | |
| rrf_v = w_vec / (k_rrf + rv) if rv is not None else 0.0 | |
| fused = rrf_b + rrf_v | |
| cand = Candidate( | |
| assessment_id=aid, | |
| source="union", | |
| bm25_rank=rb, | |
| vector_rank=rv, | |
| hybrid_rank=None, # filled after sorting by fused score | |
| bm25_score=bm25_scores.get(aid), | |
| vector_score=vec_scores.get(aid), | |
| score=fused, | |
| ) | |
| scored.append((fused, cand)) | |
| scored.sort(key=lambda x: x[0], reverse=True) | |
| for rank, (_, cand) in enumerate(scored[:topn]): | |
| cand.hybrid_rank = rank + 1 | |
| candidates.append(cand) | |
| return CandidateSet( | |
| candidates=candidates, | |
| raw_bm25=bm25_ids, | |
| raw_vector=vec_ids, | |
| fusion={"w_bm25": w_bm25, "w_vec": w_vec, "k_rrf": k_rrf}, | |
| ) | |