620 lines
12 KiB
Python
620 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
|
|
|
|
|
|
EVALUATION_DIR = PROJECT_ROOT / "evaluation"
|
|
JSON_FILES_DIR = EVALUATION_DIR / "json_files"
|
|
|
|
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()
|