retrieval

This commit is contained in:
Ján Pták 2026-08-14 22:08:33 +02:00
parent d6d02fef4d
commit 2871bc853a
3 changed files with 1568 additions and 1085 deletions

View File

@ -46,6 +46,11 @@ EVALUATION_MODES = (
"hybrid",
)
VALID_SPLITS = (
"dev",
"test",
)
def load_questions(
path: Path,
@ -75,6 +80,8 @@ def load_questions(
dict[str, Any]
] = []
seen_ids: set[str] = set()
for index, item in enumerate(
data,
start=1,
@ -96,6 +103,10 @@ def load_questions(
"question"
)
split = item.get(
"split"
)
expected_documents = item.get(
"expected_documents"
)
@ -111,6 +122,15 @@ def load_questions(
f"Položka {index} nemá platné id"
)
if question_id in seen_ids:
raise ValueError(
f"Dataset obsahuje duplicitné id: {question_id}"
)
seen_ids.add(
question_id
)
if (
not isinstance(
question,
@ -122,6 +142,12 @@ def load_questions(
f"{question_id}: chýba otázka"
)
if split not in VALID_SPLITS:
raise ValueError(
f"{question_id}: split musí byť "
"'dev' alebo 'test'"
)
if (
not isinstance(
expected_documents,
@ -153,6 +179,45 @@ def load_questions(
return questions
def filter_questions_by_split(
questions: list[
dict[str, Any]
],
split: str,
) -> list[dict[str, Any]]:
if split == "all":
return questions
return [
item
for item in questions
if item.get("split") == split
]
def count_splits(
questions: list[
dict[str, Any]
],
) -> dict[str, int]:
counts = {
"dev": 0,
"test": 0,
}
for item in questions:
split = item.get(
"split"
)
if split in counts:
counts[
split
] += 1
return counts
def load_index_document_paths(
db_file: Path,
) -> set[str]:
@ -497,11 +562,20 @@ def evaluate_question_mode(
)
row: dict[str, Any] = {
"id": question["id"],
"id": question[
"id"
],
"split": question.get(
"split"
),
"category": question.get(
"category",
"unknown",
),
"difficulty": question.get(
"difficulty",
"unknown",
),
"question": question[
"question"
],
@ -515,7 +589,9 @@ def evaluate_question_mode(
6,
),
"top_documents": [
item["document_path"]
item[
"document_path"
]
for item in ranked_documents
],
"top_source_urls": [
@ -695,11 +771,57 @@ def aggregate_by_category(
}
def aggregate_by_difficulty(
rows: list[
dict[str, Any]
],
) -> dict[
str,
dict[str, Any],
]:
grouped: dict[
str,
list[dict[str, Any]],
] = defaultdict(
list
)
for row in rows:
difficulty = str(
row.get(
"difficulty",
"unknown",
)
)
grouped[
difficulty
].append(
row
)
return {
difficulty: aggregate_metrics(
difficulty_rows
)
for (
difficulty,
difficulty_rows,
) in sorted(
grouped.items()
)
}
def print_summary(
summary: dict[
str,
dict[str, Any],
],
*,
split: str,
selected_count: int,
total_count: int,
) -> None:
print()
print(
@ -709,6 +831,19 @@ def print_summary(
"=" * 78
)
print(
f"Split: {split}"
)
print(
f"Questions: "
f"{selected_count}/{total_count}"
)
print(
"-" * 78
)
header = (
f"{'Mode':<10}"
f"{'Questions':>10}"
@ -786,7 +921,9 @@ def save_csv_results(
fieldnames = [
"id",
"split",
"category",
"difficulty",
"mode",
"question",
"first_relevant_rank",
@ -887,6 +1024,23 @@ def parse_args() -> argparse.Namespace:
),
)
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,
@ -943,10 +1097,60 @@ def parse_args() -> argparse.Namespace:
def main() -> None:
args = parse_args()
questions = load_questions(
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
@ -994,7 +1198,7 @@ def main() -> None:
start=1,
):
print(
f"[{index:02d}/{total:02d}] "
f"[{index:04d}/{total:04d}] "
f"{question['id']}: "
f"{question['question']}"
)
@ -1002,7 +1206,9 @@ def main() -> None:
mode_results = (
retrieve_all_modes(
args.db,
question["question"],
question[
"question"
],
limit=args.limit,
published_only=(
args.published_only
@ -1037,11 +1243,21 @@ def main() -> None:
],
] = {}
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
if row[
"mode"
] == mode
]
summary[
@ -1056,6 +1272,12 @@ def main() -> None:
mode_rows
)
by_difficulty[
mode
] = aggregate_by_difficulty(
mode_rows
)
payload = {
"configuration": {
"questions_file": str(
@ -1064,9 +1286,25 @@ def main() -> None:
"database": str(
args.db
),
"question_count": len(
"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
@ -1097,18 +1335,35 @@ def main() -> None:
),
},
"summary": summary,
"by_category": by_category,
"questions": evaluation_rows,
"by_category": (
by_category
),
"by_difficulty": (
by_difficulty
),
"questions": (
evaluation_rows
),
}
filename_suffix = (
args.split
)
json_path = (
args.output_dir
/ "retrieval_results.json"
/ (
"retrieval_results_"
f"{filename_suffix}.json"
)
)
csv_path = (
args.output_dir
/ "retrieval_results.csv"
/ (
"retrieval_results_"
f"{filename_suffix}.csv"
)
)
save_json_results(
@ -1122,15 +1377,24 @@ def main() -> None:
)
print_summary(
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}"
)

1179
scripts/search_core.py Normal file

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff