Hemprasad Badgujar
Add bm25 prefilter stage; unify candidates cache; validate reasoning quality
cd5f9d4 | """Step 05 potion broad — cosine narrowing logic + passthrough (no model load needed).""" | |
| from __future__ import annotations | |
| import numpy as np | |
| import polars as pl | |
| from config import BroadRetrievalSettings, Settings | |
| from retrieve import potion_broad_retrieval | |
| from retrieve.potion import _narrow_by_cosine | |
| def _pool() -> pl.DataFrame: | |
| return pl.DataFrame({ | |
| "candidate_id": ["c1", "c2", "c3", "c4", "c5"], | |
| "text": ["a", "b", "c", "d", "e"], | |
| }) | |
| def test_narrow_keeps_top_n_stable() -> None: | |
| scores = np.array([0.1, 0.9, 0.5, 0.8, 0.2]) | |
| kept = _narrow_by_cosine(_pool(), scores, top_n=2) | |
| assert kept["candidate_id"].to_list() == ["c2", "c4"] # 0.9, 0.8 | |
| assert kept["potion_score"].to_list() == [0.9, 0.8] | |
| def test_narrow_ties_break_on_candidate_id() -> None: | |
| scores = np.array([0.5, 0.5, 0.5, 0.1, 0.1]) | |
| kept = _narrow_by_cosine(_pool(), scores, top_n=2) # keep 2 of the three 0.5s → lowest ids | |
| assert kept["candidate_id"].to_list() == ["c1", "c2"] # deterministic id tiebreak | |
| def test_narrow_top_n_above_pool_size_keeps_all() -> None: | |
| scores = np.array([0.1, 0.9, 0.5]) | |
| kept = _narrow_by_cosine(_pool().head(3), scores, top_n=1000) | |
| assert kept.height == 3 # top_n clamped to the pool size, never errors | |
| def test_disabled_is_passthrough_with_zero_score() -> None: | |
| out = potion_broad_retrieval(_pool(), "query", Settings(broad_retrieval=BroadRetrievalSettings(enabled=False))) | |
| assert out.height == 5 and out["potion_score"].to_list() == [0.0] * 5 | |
| def test_empty_query_is_passthrough() -> None: | |
| out = potion_broad_retrieval(_pool(), " ", Settings()) # blank query → no model load, passthrough | |
| assert out.height == 5 and "potion_score" in out.columns | |