scripts/53_ablate_ranking.py

"""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()