retrieval_correction

evaluation/evaluate_retrieval.py

@@ -8,15 +8,10 @@ from pathlib import Path
 from typing import Any
 
 
-PROJECT_ROOT = Path(
-    __file__
-).resolve().parents[1]
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
 
 if str(PROJECT_ROOT) not in sys.path:
-    sys.path.insert(
-        0,
-        str(PROJECT_ROOT),
-    )
+    sys.path.insert(0, str(PROJECT_ROOT))
 
 
 from evaluation.metrics import (
@@ -37,11 +32,15 @@ from evaluation.retrieval_runner import (
 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],
-    ],
+    summary: dict[str, dict[str, Any]],
     *,
     split: str,
     selected_count: int,
@@ -51,6 +50,7 @@ def print_summary(
     print(
         "Retrieval evaluation"
     )
+
     print(
         "=" * 78
     )
@@ -104,6 +104,7 @@ def print_summary(
     print(
         "=" * 78
     )
+
     print()
 
 
@@ -134,9 +135,7 @@ def save_json_results(
 
 def save_csv_results(
     path: Path,
-    rows: list[
-        dict[str, Any]
-    ],
+    rows: list[dict[str, Any]],
 ) -> None:
     path.parent.mkdir(
         parents=True,
@@ -216,36 +215,22 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--questions",
         type=Path,
-        default=(
-            PROJECT_ROOT
-            / "evaluation"
-            / "questions.json"
-        ),
-        help=(
-            "Cesta k questions.json"
-        ),
+        default=QUESTIONS_PATH,
+        help="Cesta k questions.json",
     )
 
     parser.add_argument(
         "--db",
         type=Path,
         default=DB_FILE,
-        help=(
-            "Cesta k SQLite indexu"
-        ),
+        help="Cesta k SQLite indexu",
     )
 
     parser.add_argument(
         "--output-dir",
         type=Path,
-        default=(
-            PROJECT_ROOT
-            / "evaluation"
-            / "results"
-        ),
-        help=(
-            "Adresár pre výsledky"
-        ),
+        default=RESULTS_DIR,
+        help="Adresár pre výsledky",
     )
 
     parser.add_argument(
@@ -329,12 +314,10 @@ def main() -> None:
         all_questions
     )
 
-    questions = (
-        filter_questions_by_split(
+    questions = filter_questions_by_split(
         all_questions,
         args.split,
     )
-    )
 
     if not questions:
         raise RuntimeError(
@@ -346,33 +329,45 @@ def main() -> None:
     print(
         "Dataset"
     )
+
     print(
         "=" * 60
     )
 
     print(
-        f"Total: {len(all_questions)}"
+        f"Questions file: "
+        f"{args.questions}"
     )
 
     print(
-        f"Dev:   {split_counts['dev']}"
+        f"Total: "
+        f"{len(all_questions)}"
     )
 
     print(
-        f"Test:  {split_counts['test']}"
+        f"Dev:   "
+        f"{split_counts['dev']}"
     )
 
     print(
-        f"Selected split: {args.split}"
+        f"Test:  "
+        f"{split_counts['test']}"
     )
 
     print(
-        f"Selected questions: {len(questions)}"
+        f"Selected split: "
+        f"{args.split}"
+    )
+
+    print(
+        f"Selected questions: "
+        f"{len(questions)}"
     )
 
     print(
         "=" * 60
     )
+
     print()
 
     indexed_documents = (
@@ -427,8 +422,7 @@ def main() -> None:
             f"{question['question']}"
         )
 
-        mode_results = (
-            retrieve_all_modes(
+        mode_results = retrieve_all_modes(
             args.db,
             question[
                 "question"
@@ -441,7 +435,6 @@ def main() -> None:
                 args.max_per_document
             ),
         )
-        )
 
         for mode in EVALUATION_MODES:
             evaluation_rows.append(
@@ -510,9 +503,7 @@ def main() -> None:
             "database": str(
                 args.db
             ),
-            "split": (
-                args.split
-            ),
+            "split": args.split,
             "dataset_question_count": len(
                 all_questions
             ),
@@ -559,34 +550,24 @@ def main() -> None:
             ),
         },
         "summary": summary,
-        "by_category": (
-            by_category
-        ),
-        "by_difficulty": (
-            by_difficulty
-        ),
-        "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_"
-            f"{filename_suffix}.json"
+            f"retrieval_results_"
+            f"{args.split}.json"
         )
     )
 
     csv_path = (
         args.output_dir
         / (
-            "retrieval_results_"
-            f"{filename_suffix}.csv"
+            f"retrieval_results_"
+            f"{args.split}.csv"
         )
     )
 
@@ -616,11 +597,13 @@ def main() -> None:
     )
 
     print(
-        f"  JSON: {json_path}"
+        f"  JSON: "
+        f"{json_path}"
     )
 
     print(
-        f"  CSV:  {csv_path}"
+        f"  CSV:  "
+        f"{csv_path}"
     )
 
     if missing_expected: