433 lines
9.7 KiB
Python
433 lines
9.7 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[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()
|