retrieval

evaluation/evaluate_retrieval.py

@@ -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,
+        "by_category": (
-        "questions": evaluation_rows,
+            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}"
     )