"""Compare ranking controls without rebuilding indexes or calling an LLM.""" from __future__ import annotations import argparse import json import sys from datetime import datetime, timezone from pathlib import Path sys.path.insert(0, str(Path(__file__).parent.parent)) from config import DATA_DIR, DEFAULT_APP_VERSION import core.search_engine as search_module from core.search_engine import MenuSearchEngine from core.version_paths import raw_menus_path def _answers(item: dict) -> set[str]: values = item.get("expected_ids") if values is None and item.get("menu_id"): values = [item["menu_id"]] return {str(value).strip() for value in (values or []) if str(value).strip()} def _rank(predicted: list[str], expected: set[str]) -> int | None: return next( ( index for index, menu_id in enumerate(predicted, 1) if menu_id in expected ), None, ) def _metrics(ranks: list[int | None]) -> dict[str, float]: total = len(ranks) return { "n": total, "acc1": sum(rank is not None and rank <= 1 for rank in ranks) / total, "acc3": sum(rank is not None and rank <= 3 for rank in ranks) / total, "acc5": sum(rank is not None and rank <= 5 for rank in ranks) / total, "mrr5": sum(1.0 / rank for rank in ranks if rank) / total, } def _stage1_matrix() -> list[dict]: return [ { "name": "baseline", "query_expansion": True, "popularity": 0.005, "query_top_k": 50, "query_threshold": 0.85, "query_weight": 0.4, }, { "name": "expansion_off", "query_expansion": False, "popularity": 0.005, "query_top_k": 50, "query_threshold": 0.85, "query_weight": 0.4, }, { "name": "popularity_off", "query_expansion": True, "popularity": 0.0, "query_top_k": 50, "query_threshold": 0.85, "query_weight": 0.4, }, { "name": "expansion_popularity_off", "query_expansion": False, "popularity": 0.0, "query_top_k": 50, "query_threshold": 0.85, "query_weight": 0.4, }, ] def _query_index_matrix( *, query_expansion: bool, popularity: float, ) -> list[dict]: common = { "query_expansion": query_expansion, "popularity": popularity, } return [ { **common, "name": "qidx_control", "query_top_k": 50, "query_threshold": 0.85, "query_weight": 0.4, }, { **common, "name": "qidx_top100", "query_top_k": 100, "query_threshold": 0.85, "query_weight": 0.4, }, { **common, "name": "qidx_top200", "query_top_k": 200, "query_threshold": 0.85, "query_weight": 0.4, }, { **common, "name": "qidx_top50_threshold082", "query_top_k": 50, "query_threshold": 0.82, "query_weight": 0.4, }, { **common, "name": "qidx_top100_threshold082", "query_top_k": 100, "query_threshold": 0.82, "query_weight": 0.4, }, { **common, "name": "qidx_top200_threshold082", "query_top_k": 200, "query_threshold": 0.82, "query_weight": 0.4, }, { **common, "name": "qidx_top50_threshold080", "query_top_k": 50, "query_threshold": 0.80, "query_weight": 0.4, }, { **common, "name": "qidx_top100_threshold080", "query_top_k": 100, "query_threshold": 0.80, "query_weight": 0.4, }, { **common, "name": "qidx_top200_threshold080", "query_top_k": 200, "query_threshold": 0.80, "query_weight": 0.4, }, { **common, "name": "qidx_top200_weight06", "query_top_k": 200, "query_threshold": 0.85, "query_weight": 0.6, }, ] def _popularity_matrix(*, query_expansion: bool) -> list[dict]: return [ { "name": f"popularity_{weight:g}", "query_expansion": query_expansion, "popularity": weight, "query_top_k": 50, "query_threshold": 0.85, "query_weight": 0.4, } for weight in (0.005, 0.004, 0.003, 0.002, 0.001, 0.0) ] def _apply(config: dict) -> None: search_module.QUERY_EXPANSION_ENABLED = config["query_expansion"] search_module.SEARCH_COUNT_BOOST_WEIGHT = config["popularity"] search_module.QUERY_INDEX_TOP_K = config["query_top_k"] search_module.QUERY_IDX_THRESHOLD = config["query_threshold"] search_module.W_QUERY_IDX = config["query_weight"] def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--app-version", default=DEFAULT_APP_VERSION) parser.add_argument("--eval-path", type=Path, default=DATA_DIR / "eval_queries_dev.json") parser.add_argument( "--matrix", choices=("stage1", "popularity", "qidx"), default="stage1", ) parser.add_argument( "--expansion", choices=("on", "off"), default="on", help="base query-expansion setting for the qidx matrix", ) parser.add_argument( "--popularity", type=float, default=0.005, help="base search-count boost for the qidx matrix", ) parser.add_argument("--output", type=Path) args = parser.parse_args() rows = json.loads(args.eval_path.read_text(encoding="utf-8")) available = { item["menu_id"] for item in json.loads( raw_menus_path(args.app_version).read_text(encoding="utf-8") ) } rows = [ item for item in rows if not item.get("no_match", False) and _answers(item) & available ] if args.matrix == "stage1": matrix = _stage1_matrix() elif args.matrix == "popularity": matrix = _popularity_matrix( query_expansion=args.expansion == "on" ) else: matrix = _query_index_matrix( query_expansion=args.expansion == "on", popularity=args.popularity, ) engine = MenuSearchEngine.get_instance(args.app_version) original_embed = engine.embedder.embed_text embedding_cache: dict[str, list] = {} def cached_embed(text: str) -> list: if text not in embedding_cache: embedding_cache[text] = original_embed(text) return embedding_cache[text] engine.embedder.embed_text = cached_embed reports = [] baseline_ranks: list[int | None] | None = None for config in matrix: _apply(config) ranks = [] predictions = [] by_source: dict[str, list[int | None]] = {} for item in rows: hits = engine.search( item["query"], top_n=5, threshold=0.0, use_hyde=False, use_reranker=False, ) predicted = [hit.get("menu_id", "") for hit in hits] rank = _rank(predicted, _answers(item)) ranks.append(rank) by_source.setdefault( str(item.get("source") or "unknown"), [] ).append(rank) predictions.append( { "query": item["query"], "source": str(item.get("source") or "unknown"), "expected_ids": sorted(_answers(item)), "predicted_ids": predicted, "rank": rank, } ) if baseline_ranks is None: baseline_ranks = ranks improved_top3 = sum( (before is None or before > 3) and after is not None and after <= 3 for before, after in zip(baseline_ranks, ranks) ) regressed_top3 = sum( before is not None and before <= 3 and (after is None or after > 3) for before, after in zip(baseline_ranks, ranks) ) report = { "config": config, "metrics": _metrics(ranks), "metrics_by_source": { source: _metrics(source_ranks) for source, source_ranks in sorted(by_source.items()) }, "top3_improved_vs_control": improved_top3, "top3_regressed_vs_control": regressed_top3, "predictions": predictions, } reports.append(report) metrics = report["metrics"] print( f"{config['name']}: Acc@1={metrics['acc1']:.1%} " f"Acc@3={metrics['acc3']:.1%} Acc@5={metrics['acc5']:.1%} " f"MRR@5={metrics['mrr5']:.3f} " f"Top3 +{improved_top3}/-{regressed_top3}" ) for source, source_metrics in report["metrics_by_source"].items(): print( f" {source}: Acc@1={source_metrics['acc1']:.1%} " f"Acc@3={source_metrics['acc3']:.1%} " f"Acc@5={source_metrics['acc5']:.1%}" ) output = args.output or ( DATA_DIR / "experiments" / f"ranking_{args.matrix}_{args.app_version.replace('.', '_')}.json" ) output.parent.mkdir(parents=True, exist_ok=True) output.write_text( json.dumps( { "created_at": datetime.now(timezone.utc).isoformat(), "app_version": args.app_version, "eval_path": str(args.eval_path), "rows": len(rows), "reports": reports, }, ensure_ascii=False, indent=2, ) + "\n", encoding="utf-8", ) print(f"output: {output}") if __name__ == "__main__": main()