Spaces:
Runtime error
Runtime error
| from __future__ import annotations | |
| import argparse | |
| import sys | |
| from pathlib import Path | |
| ROOT = Path(__file__).resolve().parents[1] | |
| sys.path.insert(0, str(ROOT / "src")) | |
| from legal_rag.data import load_corpus, load_rag_eval, write_json | |
| from legal_rag.metrics import ndcg_at_k, recall_at_k, reciprocal_rank | |
| from legal_rag.retrievers import BM25Retriever, ChromaDBRetriever, HybridRetriever | |
| def build_retriever(args: argparse.Namespace, corpus): | |
| if args.retriever == "bm25": | |
| return BM25Retriever(corpus) | |
| collection_name = "baseline_rag_db" | |
| if args.embedding_model and ("triplet" in args.embedding_model.lower() or "finetuned" in args.embedding_model.lower()): | |
| collection_name = "finetuned_rag_db" | |
| dense = ChromaDBRetriever( | |
| corpus, | |
| model_name=args.embedding_model, | |
| persist_dir=args.index_dir.parent / "chroma_db", | |
| collection_name=collection_name, | |
| batch_size=args.batch_size, | |
| ) | |
| if args.retriever == "dense": | |
| return dense | |
| bm25 = BM25Retriever(corpus) | |
| return HybridRetriever(dense=dense, bm25=bm25, dense_weight=args.dense_weight, candidate_k=args.candidate_k) | |
| def main() -> None: | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--data-dir", type=Path, default=Path("data")) | |
| parser.add_argument("--output", type=Path, default=Path("outputs/retrieval_eval.json")) | |
| parser.add_argument("--index-dir", type=Path, default=Path("outputs/index")) | |
| parser.add_argument("--retriever", choices=["dense", "bm25", "hybrid"], default="dense") | |
| parser.add_argument("--embedding-model", default="intfloat/multilingual-e5-base") | |
| parser.add_argument("--top-k", type=int, default=10) | |
| parser.add_argument("--candidate-k", type=int, default=50) | |
| parser.add_argument("--dense-weight", type=float, default=0.5) | |
| parser.add_argument("--batch-size", type=int, default=64) | |
| parser.add_argument("--limit", type=int, default=None) | |
| args = parser.parse_args() | |
| corpus = load_corpus(args.data_dir) | |
| eval_rows = load_rag_eval(args.data_dir) | |
| if args.limit: | |
| eval_rows = eval_rows[: args.limit] | |
| retriever = build_retriever(args, corpus) | |
| per_query = [] | |
| totals = {"recall@5": 0.0, "recall@10": 0.0, "mrr": 0.0, "ndcg@10": 0.0} | |
| queries = [row["query"] for row in eval_rows] | |
| result_batches = retriever.batch_search(queries, top_k=max(args.top_k, 10)) | |
| for row, results in zip(eval_rows, result_batches): | |
| retrieved_ids = [result.doc.id for result in results] | |
| gold_ids = set(row["gold_chunk_ids"]) | |
| metrics = { | |
| "recall@5": recall_at_k(retrieved_ids, gold_ids, 5), | |
| "recall@10": recall_at_k(retrieved_ids, gold_ids, 10), | |
| "mrr": reciprocal_rank(retrieved_ids, gold_ids), | |
| "ndcg@10": ndcg_at_k(retrieved_ids, gold_ids, 10), | |
| } | |
| for key, value in metrics.items(): | |
| totals[key] += value | |
| per_query.append( | |
| { | |
| "query_id": row["query_id"], | |
| "query": row["query"], | |
| "gold_chunk_ids": row["gold_chunk_ids"], | |
| "retrieved_ids": retrieved_ids[: args.top_k], | |
| "metrics": metrics, | |
| } | |
| ) | |
| summary = {key: value / len(eval_rows) for key, value in totals.items()} | |
| output = { | |
| "config": { | |
| "retriever": args.retriever, | |
| "embedding_model": args.embedding_model if args.retriever != "bm25" else None, | |
| "top_k": args.top_k, | |
| "dense_weight": args.dense_weight if args.retriever == "hybrid" else None, | |
| "num_queries": len(eval_rows), | |
| }, | |
| "summary": summary, | |
| "per_query": per_query, | |
| } | |
| write_json(args.output, output) | |
| print("Retrieval evaluation complete") | |
| print(output["config"]) | |
| print(summary) | |
| print(f"Wrote {args.output}") | |
| if __name__ == "__main__": | |
| main() | |