expose reading from file in run_rankllm (#320)

* add support for passing a retrieval file (Cached_file retrieval mode)
* add output files and qrel files to args, add populate invocations history to args to enable writing results in cached_file mode.

Author: sahel-sh

src/rank_llm/retrieve/retriever.py

@@ -17,6 +17,7 @@
 class RetrievalMode(Enum):
     DATASET = "dataset"
     CUSTOM = "custom"
+    CACHED_FILE = "cached_file"
 
     def __str__(self):
         return self.value

src/rank_llm/retrieve_and_rerank.py

@@ -1,7 +1,8 @@
 import copy
+from pathlib import Path
 from typing import Any, Dict, List, Optional, Union
 
-from rank_llm.data import Query, Request
+from rank_llm.data import DataWriter, Query, Request, read_requests_from_file
 from rank_llm.rerank import IdentityReranker, RankLLM, Reranker
 from rank_llm.rerank.reranker import extract_kwargs
 from rank_llm.retrieve import (
@@ -79,7 +80,6 @@ def retrieve_and_rerank(
         # Reranker is of type RankLLM
         for pass_ct in range(num_passes):
             print(f"Pass {pass_ct + 1} of {num_passes}:")
-
             rerank_results = reranker.rerank_batch(
                 requests,
                 rank_end=top_k_retrieve,
@@ -125,6 +125,55 @@ def retrieve_and_rerank(
             EvalFunction.eval(["-c", "-m", "ndcg_cut.10", TOPICS[dataset], file_name])
         else:
             print(f"Skipping evaluation as {dataset} is not in TOPICS.")
+    elif (
+        retrieval_mode == RetrievalMode.CACHED_FILE
+        and reranker.get_model_coordinator() is not None
+    ):
+        writer = DataWriter(rerank_results)
+        keys_and_defaults = [
+            ("output_jsonl_file", ""),
+            ("output_trec_file", ""),
+            ("invocations_history_file", ""),
+        ]
+        [
+            output_jsonl_file,
+            output_trec_file,
+            invocations_history_file,
+        ] = extract_kwargs(keys_and_defaults, **kwargs)
+        if output_jsonl_file:
+            path = Path(output_jsonl_file)
+            path.parent.mkdir(parents=True, exist_ok=True)
+            writer.write_in_jsonl_format(output_jsonl_file)
+        if output_trec_file:
+            path = Path(output_trec_file)
+            path.parent.mkdir(parents=True, exist_ok=True)
+            writer.write_in_trec_eval_format(output_trec_file)
+        keys_and_defaults = [("populate_invocations_history", False)]
+        [populate_invocations_history] = extract_kwargs(keys_and_defaults, **kwargs)
+        if populate_invocations_history:
+            if invocations_history_file:
+                path = Path(invocations_history_file)
+                path.parent.mkdir(parents=True, exist_ok=True)
+                writer.write_inference_invocations_history(invocations_history_file)
+            else:
+                raise ValueError(
+                    "--invocations_history_file must be a valid jsonl file to store invocations history."
+                )
+        keys_and_defaults = [("qrels_file", "")]
+        [qrels_file] = extract_kwargs(keys_and_defaults, **kwargs)
+        if qrels_file:
+            from rank_llm.evaluation.trec_eval import EvalFunction
+
+            print("Evaluating:")
+            EvalFunction.from_results(
+                rerank_results, qrels_file, ["-c", "-m", "ndcg_cut.1"]
+            )
+            EvalFunction.from_results(
+                rerank_results, qrels_file, ["-c", "-m", "ndcg_cut.5"]
+            )
+            EvalFunction.from_results(
+                rerank_results, qrels_file, ["-c", "-m", "ndcg_cut.10"]
+            )
 
     if interactive:
         return (rerank_results, reranker.get_model_coordinator())
@@ -211,5 +260,11 @@ def retrieve(
         requests = Retriever.from_custom_index(
             index_path=index_path, topics_path=topics_path, index_type=index_type
         )
+    elif retrieval_mode == RetrievalMode.CACHED_FILE:
+        keys_and_defaults = [
+            ("requests_file", ""),
+        ]
+        [requests_file] = extract_kwargs(keys_and_defaults, **kwargs)
+        requests = read_requests_from_file(requests_file)
 
     return requests

src/rank_llm/scripts/run_rank_llm.py

@@ -28,6 +28,11 @@ def main(args):
     dataset = args.dataset
     num_gpus = args.num_gpus
     retrieval_method = args.retrieval_method
+    requests_file = args.requests_file
+    qrels_file = args.qrels_file
+    output_jsonl_file = args.output_jsonl_file
+    output_trec_file = args.output_trec_file
+    invocations_history_file = args.invocations_history_file
     prompt_template_path = args.prompt_template_path
     num_few_shot_examples = args.num_few_shot_examples
     few_shot_file = args.few_shot_file
@@ -36,7 +41,9 @@ def main(args):
     num_few_shot_examples = args.num_few_shot_examples
     device = "cuda" if torch.cuda.is_available() else "cpu"
     variable_passages = args.variable_passages
-    retrieval_mode = RetrievalMode.DATASET
+    retrieval_mode = (
+        RetrievalMode.DATASET if args.dataset else RetrievalMode.CACHED_FILE
+    )
     num_passes = args.num_passes
     stride = args.stride
     window_size = args.window_size
@@ -49,12 +56,26 @@ def main(args):
     sglang_batched = args.sglang_batched
     tensorrt_batched = args.tensorrt_batched
 
+    if args.requests_file:
+        if args.retrieval_method:
+            parser.error("--retrieval_method must not be used with --requests_file")
+        if not os.path.exists(args.requests_file):
+            parser.error(f"--requests_file not found: {args.requests_file}")
+
+    if args.dataset and not args.retrieval_method:
+        parser.error("--retrieval_method is required when --dataset is provided")
+
     _ = retrieve_and_rerank(
         model_path=model_path,
         query=query,
         batch_size=batch_size,
         dataset=dataset,
         retrieval_mode=retrieval_mode,
+        requests_file=requests_file,
+        qrels_file=qrels_file,
+        output_jsonl_file=output_jsonl_file,
+        output_trec_file=output_trec_file,
+        invocations_history_file=invocations_history_file,
         retrieval_method=retrieval_method,
         top_k_retrieve=top_k_candidates,
         top_k_rerank=top_k_rerank,
@@ -142,21 +163,46 @@ def main(args):
         default=None,
         help="the max number of queries to process from the dataset",
     )
-    parser.add_argument(
+    retrieval_input_group = parser.add_mutually_exclusive_group(required=True)
+    retrieval_input_group.add_argument(
         "--dataset",
         type=str,
-        required=True,
-        help=f"Should be one of 1- dataset name, must be in {TOPICS.keys()},  2- a list of inline documents  3- a list of inline hits 4- filename containing retrieved results",
-    )
-    parser.add_argument(
-        "--num_gpus", type=int, default=1, help="the number of GPUs to use"
+        help=f"Should be one of 1- dataset name, must be in {TOPICS.keys()},  2- a list of inline documents  3- a list of inline hits; must be used when --requests_file is not specified",
     )
     parser.add_argument(
         "--retrieval_method",
         type=RetrievalMethod,
-        required=True,
+        help="Required if --dataset is used; must be omitted with --requests_file",
         choices=list(RetrievalMethod),
     )
+    retrieval_input_group.add_argument(
+        "--requests_file",
+        type=str,
+        help=f"Path to a JSONL file containing requests; must be used when --dataset is not specified.",
+    )
+    parser.add_argument(
+        "--qrels_file",
+        type=str,
+        help="Only used with --requests_file; when present the Trec eval will be executed using this qrels file",
+    )
+    parser.add_argument(
+        "--output_jsonl_file",
+        type=str,
+        help="Only used with --requests_file; when present, the ranked results will be saved in this JSONL file.",
+    )
+    parser.add_argument(
+        "--output_trec_file",
+        type=str,
+        help="Only used with --requests_file; when present, the ranked results will be saved in this txt file in trec format.",
+    )
+    parser.add_argument(
+        "--invocations_history_file",
+        type=str,
+        help="Only used with --requests_file and --populate_invocations_history; when present, the LLM invocations history (prompts, completions, and input/output token counts) will be stored in this file.",
+    )
+    parser.add_argument(
+        "--num_gpus", type=int, default=1, help="the number of GPUs to use"
+    )
     parser.add_argument(
         "--prompt_mode",
         type=PromptMode,