Updated RankLLM integration into rerankers (#67)

* updated rankllm implementation along with new arguments

* small change

* updated integration with new fields

* updated integration with new fields

Author: clides

rerankers/models/rankllm_ranker.py

@@ -4,34 +4,69 @@
 from rerankers.results import RankedResults, Result
 from rerankers.utils import prep_docs
 
-from rank_llm.data import Candidate, Query, Request
-from rank_llm.rerank.listwise.vicuna_reranker import VicunaReranker
-from rank_llm.rerank.listwise.zephyr_reranker import ZephyrReranker
-from rank_llm.rerank.listwise.rank_gpt import SafeOpenai
+# from rerankers import Reranker
+
 from rank_llm.rerank.reranker import Reranker as rankllm_Reranker
+from rank_llm.rerank import PromptMode, get_azure_openai_args, get_genai_api_key, get_openai_api_key
+from rank_llm.data import Candidate, Query, Request
 
 
 class RankLLMRanker(BaseRanker):
     def __init__(
         self,
-        model: str,
+        model: str = "rank_zephyr",
         api_key: Optional[str] = None,
         lang: str = "en",
         verbose: int = 1,
+        # RankLLM specific arguments
+        window_size: int = 20,
+        context_size: int = 4096,
+        prompt_mode: PromptMode = PromptMode.RANK_GPT,
+        num_few_shot_examples: int = 0,
+        few_shot_file: Optional[str] = None,
+        num_gpus: int = 1,
+        variable_passages: bool = False,
+        use_logits: bool = False,
+        use_alpha: bool = False,
+        stride: int = 10,
+        use_azure_openai: bool = False,
     ) -> "RankLLMRanker":
         self.api_key = api_key
         self.model = model
         self.verbose = verbose
         self.lang = lang
-
-        if "zephyr" in self.model.lower():
-            self.rankllm_ranker = ZephyrReranker()
-        elif "vicuna" in self.model.lower():
-            self.rankllm_ranker = VicunaReranker()
-        elif "gpt" in self.model.lower():
-            self.rankllm_ranker = rankllm_Reranker(
-                SafeOpenai(model=self.model, context_size=4096, keys=self.api_key)
-            )
+        
+        # RankLLM-specific parameters
+        self.window_size = window_size
+        self.context_size = context_size
+        self.prompt_mode = prompt_mode
+        self.num_few_shot_examples = num_few_shot_examples
+        self.few_shot_file = few_shot_file
+        self.num_gpus = num_gpus
+        self.variable_passages = variable_passages
+        self.use_logits = use_logits
+        self.use_alpha = use_alpha
+        self.stride = stride
+        self.use_azure_openai = use_azure_openai
+        
+        kwargs = {
+            "model_path": self.model,
+            "default_model_coordinator": None,
+            "context_size": self.context_size,
+            "prompt_mode": self.prompt_mode,
+            "num_gpus": self.num_gpus,
+            "use_logits": self.use_logits,
+            "use_alpha": self.use_alpha,
+            "num_few_shot_examples": self.num_few_shot_examples,
+            "few_shot_file": self.few_shot_file,
+            "variable_passages": self.variable_passages,
+            "interactive": False,
+            "window_size": self.window_size,
+            "stride": self.stride,
+            "use_azure_openai": self.use_azure_openai,
+        }
+        model_coordinator = rankllm_Reranker.create_model_coordinator(**kwargs)
+        self.reranker = rankllm_Reranker(model_coordinator)
 
     def rank(
         self,
@@ -52,7 +87,7 @@ def rank(
             ],
         )
 
-        rankllm_results = self.rankllm_ranker.rerank(
+        rankllm_results = self.reranker.rerank(
             request,
             rank_end=len(docs) if rank_end == 0 else rank_end,
             window_size=min(20, len(docs)),

rerankers/reranker.py

@@ -22,7 +22,7 @@
     "rankgpt": {"en": "gpt-4-turbo-preview", "other": "gpt-4-turbo-preview"},
     "rankgpt3": {"en": "gpt-3.5-turbo", "other": "gpt-3.5-turbo"},
     "rankgpt4": {"en": "gpt-4", "other": "gpt-4"},
-    "rankllm": {"en": "gpt-4o", "other": "gpt-4o"},
+    "rankllm": {"en": "rank_zephyr", "other": "rank_zephyr"},
     "colbert": {
         "en": "colbert-ir/colbertv2.0",
         "fr": "bclavie/FraColBERTv2",