dp-zp-agent/evaluation/evaluate_retrieval.py
2026-08-15 01:00:09 +02:00

637 lines
12 KiB
Python

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[1]
if str(PROJECT_ROOT) not in sys.path:
sys.path.insert(
0,
str(PROJECT_ROOT),
)
from evaluation.metrics import (
aggregate_by_category,
aggregate_by_difficulty,
aggregate_metrics,
evaluate_question_mode,
)
from evaluation.retrieval_runner import (
EVALUATION_MODES,
count_splits,
filter_questions_by_split,
load_index_document_paths,
load_questions,
retrieve_all_modes,
validate_dataset,
)
from scripts.common import DB_FILE
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=(
PROJECT_ROOT
/ "evaluation"
/ "questions.json"
),
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=(
PROJECT_ROOT
/ "evaluation"
/ "results"
),
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"Total: {len(all_questions)}"
)
print(
f"Dev: {split_counts['dev']}"
)
print(
f"Test: {split_counts['test']}"
)
print(
f"Selected split: {args.split}"
)
print(
f"Selected questions: {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
),
}
filename_suffix = (
args.split
)
json_path = (
args.output_dir
/ (
"retrieval_results_"
f"{filename_suffix}.json"
)
)
csv_path = (
args.output_dir
/ (
"retrieval_results_"
f"{filename_suffix}.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: {json_path}"
)
print(
f" CSV: {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()