from __future__ import annotations import argparse import csv import json import sys from pathlib import Path from typing import Any PROJECT_ROOT = Path(__file__).resolve().parents[3] if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) from tools.rag.evaluation.metrics import ( aggregate_by_category, aggregate_by_difficulty, aggregate_metrics, evaluate_question_mode, ) from tools.rag.evaluation.retrieval_runner import ( EVALUATION_MODES, count_splits, filter_questions_by_split, load_index_document_paths, load_questions, retrieve_all_modes, validate_dataset, ) from tools.rag.common import DB_FILE EVALUATION_DIR = PROJECT_ROOT / "experiments/rag" JSON_FILES_DIR = EVALUATION_DIR / "datasets" QUESTIONS_PATH = JSON_FILES_DIR / "questions.json" RESULTS_DIR = EVALUATION_DIR / "results" def print_summary( summary: dict[str, dict[str, Any]], *, split: str, selected_count: int, total_count: int, ) -> None: print() print("Retrieval evaluation") print("=" * 78) print(f"Split: {split}") print(f"Questions: " f"{selected_count}/{total_count}") print("-" * 78) header = ( f"{'Mode':<10}" f"{'Questions':>10}" f"{'Hit@1':>10}" f"{'Hit@3':>10}" f"{'Hit@5':>10}" f"{'MRR':>10}" f"{'Recall@5':>12}" ) print(header) print("-" * 78) for mode in EVALUATION_MODES: metrics = summary[mode] print( f"{mode:<10}" f"{metrics['questions']:>10}" f"{metrics['hit_at_1']:>10.3f}" f"{metrics['hit_at_3']:>10.3f}" f"{metrics['hit_at_5']:>10.3f}" f"{metrics['mrr']:>10.3f}" f"{metrics['recall_at_5']:>12.3f}" ) print("=" * 78) print() def save_json_results( path: Path, payload: dict[str, Any], ) -> None: path.parent.mkdir( parents=True, exist_ok=True, ) with path.open( "w", encoding="utf-8", ) as file: json.dump( payload, file, ensure_ascii=False, indent=2, ) file.write("\n") def save_csv_results( path: Path, rows: list[dict[str, Any]], ) -> None: path.parent.mkdir( parents=True, exist_ok=True, ) fieldnames = [ "id", "split", "category", "difficulty", "mode", "question", "first_relevant_rank", "reciprocal_rank", "hit_at_1", "hit_at_3", "hit_at_5", "recall_at_1", "recall_at_3", "recall_at_5", "expected_documents", "top_documents", "top_source_urls", ] with path.open( "w", encoding="utf-8", newline="", ) as file: writer = csv.DictWriter( file, fieldnames=fieldnames, ) writer.writeheader() for row in rows: csv_row = dict(row) for key in ( "expected_documents", "top_documents", "top_source_urls", ): csv_row[key] = json.dumps( csv_row.get( key, [], ), ensure_ascii=False, ) writer.writerow({key: csv_row.get(key) for key in fieldnames}) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description=("Porovnanie FTS5, embeddingového " "a hybridného retrievalu.") ) parser.add_argument( "--questions", type=Path, default=QUESTIONS_PATH, help="Cesta k questions.json", ) parser.add_argument( "--db", type=Path, default=DB_FILE, help="Cesta k SQLite indexu", ) parser.add_argument( "--output-dir", type=Path, default=RESULTS_DIR, help="Adresár pre výsledky", ) parser.add_argument( "--split", choices=( "dev", "test", "all", ), default="dev", help=( "Časť datasetu: " "dev = ladenie, " "test = finálne hodnotenie, " "all = celý dataset. " "Predvolené je dev." ), ) parser.add_argument( "--limit", type=int, default=5, help=("Počet dokumentov použitých " "pri evaluácii. Minimum je 5."), ) parser.add_argument( "--published-only", action="store_true", help=("Vyhodnocovať iba publikované dokumenty"), ) parser.add_argument( "--max-per-document", type=int, default=1, help=("Maximálny počet chunkov " "z jedného dokumentu"), ) parser.add_argument( "--strict-dataset", action="store_true", help=("Ukončiť evaluáciu chybou, " "ak expected document nie je v indexe."), ) args = parser.parse_args() if args.limit < 5: parser.error("--limit musí byť aspoň 5, " "pretože meriame Hit@5") if args.max_per_document < 0: parser.error("--max-per-document nesmie byť záporné") return args def main() -> None: args = parse_args() all_questions = load_questions(args.questions) split_counts = count_splits(all_questions) questions = filter_questions_by_split( all_questions, args.split, ) if not questions: raise RuntimeError(f"Pre split '{args.split}' " "sa nenašli žiadne otázky.") print() print("Dataset") print("=" * 60) print(f"Questions file: " f"{args.questions}") print(f"Total: " f"{len(all_questions)}") print(f"Dev: " f"{split_counts['dev']}") print(f"Test: " f"{split_counts['test']}") print(f"Selected split: " f"{args.split}") print(f"Selected questions: " f"{len(questions)}") print("=" * 60) print() indexed_documents = load_index_document_paths(args.db) missing_expected = validate_dataset( questions, indexed_documents, ) if missing_expected: print() print("POZOR: niektoré expected_documents " "sa nenachádzajú v indexe:") for item in missing_expected: print(f" {item['question_id']}: " f"{item['document_path']}") print() if args.strict_dataset: raise RuntimeError("Evaluačný dataset obsahuje " "neexistujúce expected_documents.") evaluation_rows: list[dict[str, Any]] = [] total = len(questions) for index, question in enumerate( questions, start=1, ): print(f"[{index:04d}/{total:04d}] " f"{question['id']}: " f"{question['question']}") mode_results = retrieve_all_modes( args.db, question["question"], limit=args.limit, published_only=(args.published_only), max_per_document=(args.max_per_document), ) for mode in EVALUATION_MODES: evaluation_rows.append( evaluate_question_mode( question, mode, mode_results[mode], ) ) summary: dict[ str, dict[str, Any], ] = {} by_category: dict[ str, dict[ str, dict[str, Any], ], ] = {} by_difficulty: dict[ str, dict[ str, dict[str, Any], ], ] = {} for mode in EVALUATION_MODES: mode_rows = [row for row in evaluation_rows if row["mode"] == mode] summary[mode] = aggregate_metrics(mode_rows) by_category[mode] = aggregate_by_category(mode_rows) by_difficulty[mode] = aggregate_by_difficulty(mode_rows) payload = { "configuration": { "questions_file": str(args.questions), "database": str(args.db), "split": args.split, "dataset_question_count": len(all_questions), "selected_question_count": len(questions), "dev_question_count": (split_counts["dev"]), "test_question_count": (split_counts["test"]), "limit": args.limit, "published_only": (args.published_only), "max_per_document": (args.max_per_document), "modes": list(EVALUATION_MODES), "metrics": [ "Hit@1", "Hit@3", "Hit@5", "MRR", "Recall@5", ], }, "dataset_validation": { "indexed_document_count": len(indexed_documents), "missing_expected_document_count": len(missing_expected), "missing_expected_documents": (missing_expected), }, "summary": summary, "by_category": by_category, "by_difficulty": by_difficulty, "questions": evaluation_rows, } json_path = args.output_dir / (f"retrieval_results_" f"{args.split}.json") csv_path = args.output_dir / (f"retrieval_results_" f"{args.split}.csv") save_json_results( json_path, payload, ) save_csv_results( csv_path, evaluation_rows, ) print_summary( summary, split=args.split, selected_count=len(questions), total_count=len(all_questions), ) print("Výsledky:") print(f" JSON: " f"{json_path}") print(f" CSV: " f"{csv_path}") if missing_expected: print() print( "POZOR: pred použitím metrík " "v diplomovej práci oprav " "missing expected documents." ) if __name__ == "__main__": main()