McQbis/document-intelligence-rag

1import argparse2import sys3from pathlib import Path4 5# allow running script without installing package6sys.path.insert(0, str(Path(__file__).resolve().parents[1]))7 8from rag.retrieval.embeddings import EmbeddingModel9from rag.retrieval.retriever import HybridRetriever10from rag.routing.router import QueryRouter, RouteMode11from rag.evaluation.beir_eval import BEIREvaluator12 13 14def parse_args():15    p = argparse.ArgumentParser(description="BEIR offline evaluation")16    p.add_argument("--dataset", default="fiqa", help="BEIR dataset name")17    p.add_argument("--split", default="test")18    p.add_argument("--top-k", type=int, default=10)19    p.add_argument("--candidate-k", type=int, default=30)20    p.add_argument("--max-queries", type=int, default=None)21    p.add_argument("--model", default="BAAI/bge-base-en-v1.5", help="Embedding model or local path")22    p.add_argument("--reranker", default="BAAI/bge-reranker-base")23    p.add_argument("--data-dir", default="beir-data")24    p.add_argument(25        "--router",26        default=None,27        choices=["auto", "fast", "deep"],28        help="Użyj QueryRouter zamiast HybridRetriever. "29             "auto=heurystyka, fast=bez rerankera, deep=zawsze reranker",30    )31    return p.parse_args()32 33 34def build_retriever(args):35    emb = EmbeddingModel(model_name=args.model)36    retriever = HybridRetriever(emb, reranker_model=args.reranker)37 38    if args.router is None:39        print(f"[eval] Retriever   : HybridRetriever (rerank=True)")40        return retriever41 42    mode_map = {"auto": RouteMode.AUTO, "fast": RouteMode.FAST, "deep": RouteMode.DEEP}43    forced_mode = mode_map[args.router]44    router = QueryRouter(45        retriever,46        top_k=args.top_k,47        candidate_k=args.candidate_k,48    )49    print(f"[eval] Retriever   : QueryRouter (mode={args.router})")50 51    class _ForcedRouter:52        # Adapter forcing a fixed routing strategy for evaluation53        def __init__(self, router, mode):54            self._router = router55            self._mode = mode56 57        def build_index(self, chunks):58            self._router.retriever.build_index(chunks)59 60        def search(self, query, top_k=10, candidate_k=30):61            return self._router.search(query, mode=self._mode, top_k=top_k, candidate_k=candidate_k)62 63    return _ForcedRouter(router, forced_mode)64 65 66def main():67    args = parse_args()68 69    print(f"[eval] Model       : {args.model}")70    print(f"[eval] Reranker    : {args.reranker}")71    print(f"[eval] Dataset     : {args.dataset}/{args.split}")72    print(f"[eval] Top-K       : {args.top_k}")73    print(f"[eval] Candidate-K : {args.candidate_k}")74 75    retriever = build_retriever(args)76 77    evaluator = BEIREvaluator(retriever, data_dir=args.data_dir)78    results = evaluator.run(79        dataset=args.dataset,80        split=args.split,81        top_k=args.top_k,82        max_queries=args.max_queries,83    )84 85    print(results)86 87 88if __name__ == "__main__":89    main()