Add regression test script and max_queries arg (#246)

- added scripts for regression tests, the tests run with a select number of queries (max_queries), on a mix of datasets and retrievers.

Author: clides

regression_test.sh

@@ -0,0 +1,58 @@
+#!/bin/bash
+
+# Define test cases:
+# Format:
+#   TEST_NAMES=("Test 1" "Test 2" ...)
+#   TEST_COMMANDS=("command1" "command2" ...)
+#   TEST_EXPECTED_SCORES=(0.123 0.456 ...)
+
+TEST_NAMES=(
+  "FirstMistral (Alpha, Logits)"
+  "RZ"
+  "Qwen (Alpha)"
+  "Monot5"
+  "Duot5"
+)
+
+TEST_COMMANDS=(
+  "python src/rank_llm/scripts/run_rank_llm.py --model_path=castorini/first_mistral --top_k_candidates=50 --dataset=dl19 --retrieval_method=bm25 --prompt_mode=rank_GPT --context_size=4096 --use_alpha --use_logits --max_queries=3"
+  "python src/rank_llm/scripts/run_rank_llm.py --model_path=castorini/rank_zephyr_7b_v1_full  --top_k_candidates=50 --dataset=dl20  --retrieval_method=SPLADE++_EnsembleDistil_ONNX --prompt_mode=rank_GPT  --context_size=4096 --max_queries=3"
+  "python src/rank_llm/scripts/run_rank_llm.py  --model_path=Qwen/Qwen2.5-7B-Instruct --top_k_candidates=50 --dataset=dl21 --retrieval_method=bm25 --prompt_mode=rank_GPT  --context_size=4096 --variable_passages --max_queries=3"
+  "python src/rank_llm/scripts/run_rank_llm.py  --model_path=castorini/monot5-3b-msmarco-10k --top_k_candidates=50 --dataset=dl22 --retrieval_method=bm25 --prompt_mode=rank_GPT  --context_size=4096 --variable_passages --max_queries=3"
+   "python src/rank_llm/scripts/run_rank_llm.py  --model_path=castorini/duot5-3b-msmarco-10k --top_k_candidates=50 --dataset=dl23 --retrieval_method=bm25 --prompt_mode=rank_GPT  --context_size=4096 --variable_passages --max_queries=1"
+)
+
+TEST_EXPECTED_SCORES=(
+  0.8085
+  0.7662
+  0.7157
+  0.3997
+  0.7246
+)
+
+for i in "${!TEST_NAMES[@]}"; do
+  NAME="${TEST_NAMES[$i]}"
+  COMMAND="${TEST_COMMANDS[$i]}"
+  EXPECTED_SCORE="${TEST_EXPECTED_SCORES[$i]}"
+
+  echo "Running $NAME..."
+
+  OUTPUT=$(eval "$COMMAND" 2>&1)
+
+  SCORE=$(echo "$OUTPUT" | grep -oP 'ndcg_cut_10\s+all\s+\K\d+\.\d+')
+
+  if [ -z "$SCORE" ]; then
+    echo "❌ ERROR: Could not extract nDCG@10 score for '$NAME'"
+    continue
+  fi
+
+  LOWER_BOUND=$(echo "$EXPECTED_SCORE * 0.975" | bc -l)
+  UPPER_BOUND=$(echo "$EXPECTED_SCORE * 1.025" | bc -l)
+  PASSED=$(echo "$SCORE >= $LOWER_BOUND && $SCORE <= $UPPER_BOUND" | bc -l)
+
+  if [ "$PASSED" -eq 1 ]; then
+    echo "$NAME: PASS ✅ (Actual Score: $SCORE, Expected Score: $EXPECTED_SCORE)"
+  else
+    echo "$NAME: FAIL ❌ (Actual Score: $SCORE, Expected Score: $EXPECTED_SCORE)"
+  fi
+done

src/rank_llm/rerank/reranker.py

@@ -342,8 +342,8 @@ def create_model_coordinator(
                 ("prompt_mode", PromptMode.MONOT5),
                 (
                     "prompt_template_path",
-                    None,
-                ),  # TODO(issue #236): Need to modify and add default MONOT5 template
+                    "src/rank_llm/rerank/prompt_templates/monot5_template.yaml",
+                ),
                 ("context_size", 512),
                 ("num_few_shot_examples", 0),
                 ("few_shot_file", None),
@@ -375,7 +375,7 @@ def create_model_coordinator(
                 batch_size=batch_size,
             )
         elif "duot5" in model_path:
-            # using monot5
+            # using duot5
             print(f"Loading {model_path} ...")
 
             model_full_paths = {"duot5": "castorini/duot5-3b-msmarco-10k"}
@@ -384,8 +384,8 @@ def create_model_coordinator(
                 ("prompt_mode", PromptMode.DUOT5),
                 (
                     "prompt_template_path",
-                    None,
-                ),  # TODO(issue #236): Need to modify and add default DUOT5 template
+                    "src/rank_llm/rerank/prompt_templates/duot5_template.yaml",
+                ),
                 ("context_size", 512),
                 ("device", "cuda"),
                 ("batch_size", 64),

src/rank_llm/retrieve_and_rerank.py

@@ -1,5 +1,5 @@
 import copy
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, Optional, Union
 
 from rank_llm.data import Query, Request
 from rank_llm.rerank import IdentityReranker, RankLLM, Reranker
@@ -21,6 +21,7 @@ def retrieve_and_rerank(
     retrieval_method: RetrievalMethod = RetrievalMethod.BM25,
     top_k_retrieve: int = 50,
     top_k_rerank: int = 10,
+    max_queries: Optional[int] = None,
     shuffle_candidates: bool = False,
     print_prompts_responses: bool = False,
     qid: int = 1,
@@ -58,6 +59,9 @@ def retrieve_and_rerank(
         **kwargs,
     )
 
+    if max_queries is not None:
+        requests = requests[: min(len(requests), max_queries)]
+
     for request in requests:
         request.candidates = request.candidates[:top_k_retrieve]
 

src/rank_llm/scripts/run_rank_llm.py

@@ -22,6 +22,7 @@ def main(args):
     context_size = args.context_size
     top_k_candidates = args.top_k_candidates
     top_k_rerank = top_k_candidates if args.top_k_rerank == -1 else args.top_k_rerank
+    max_queries = args.max_queries
     dataset = args.dataset
     num_gpus = args.num_gpus
     retrieval_method = args.retrieval_method
@@ -56,6 +57,7 @@ def main(args):
         retrieval_method=retrieval_method,
         top_k_retrieve=top_k_candidates,
         top_k_rerank=top_k_rerank,
+        max_queries=max_queries,
         context_size=context_size,
         device=device,
         num_gpus=num_gpus,
@@ -119,6 +121,12 @@ def main(args):
         default=-1,
         help="the number of top candidates to return from reranking",
     )
+    parser.add_argument(
+        "--max_queries",
+        type=int,
+        default=None,
+        help="the max number of queries to process from the dataset",
+    )
     parser.add_argument(
         "--dataset",
         type=str,