reformatted

Author: Yaohui2019

src/rank_llm/rerank/listwise/__init__.py

@@ -1,16 +1,15 @@
 from .rank_gemini import SafeGenai
 from .rank_gpt import SafeOpenai
 from .rank_listwise_os_llm import RankListwiseOSLLM
+from .rankk_reranker import RankKReranker
 from .vicuna_reranker import VicunaReranker
 from .zephyr_reranker import ZephyrReranker
-from .rankk_reranker import RankKReranker
 
 __all__ = [
     "RankListwiseOSLLM",
     "VicunaReranker",
     "ZephyrReranker",
     "SafeOpenai",
     "SafeGenai",
-    "RankKReranker"
+    "RankKReranker",
 ]
-

src/rank_llm/rerank/listwise/rankk_reranker.py

@@ -1,40 +1,48 @@
 from typing import Optional
 
-from rank_llm.data import  Result
+from rank_llm.data import Result
 from rank_llm.rerank import PromptMode
 from rank_llm.rerank.listwise import RankListwiseOSLLM
 
 
 class RankKReranker(RankListwiseOSLLM):
     def __init__(
-            self,
-            model: str = "hltcoe/Rank-K-32B",
-            context_size: int = 4096,
-            prompt_mode: PromptMode = PromptMode.RANK_GPT,
-            prompt_template_path: Optional[str] = "src/rank_llm/rerank/prompt_templates/rank_k_template.yaml",
-            num_few_shot_examples: int = 0,
-            device: str = "cuda",
-            num_gpus: int = 1,
-            variable_passages: bool = True,
-            window_size: int = 20,
-            use_alpha: bool = False
+        self,
+        model: str = "hltcoe/Rank-K-32B",
+        context_size: int = 4096,
+        prompt_mode: PromptMode = PromptMode.RANK_GPT,
+        prompt_template_path: Optional[
+            str
+        ] = "src/rank_llm/rerank/prompt_templates/rank_k_template.yaml",
+        num_few_shot_examples: int = 0,
+        device: str = "cuda",
+        num_gpus: int = 1,
+        variable_passages: bool = True,
+        window_size: int = 20,
+        use_alpha: bool = False,
     ) -> None:
         super().__init__(
             model=model,
             context_size=context_size,
             prompt_mode=prompt_mode,
             prompt_template_path=prompt_template_path,
-            num_few_shot_examples = num_few_shot_examples,
-            device = device,
+            num_few_shot_examples=num_few_shot_examples,
+            device=device,
             num_gpus=num_gpus,
-            variable_passages= variable_passages,
+            variable_passages=variable_passages,
             is_thinking=True,
             reasoning_token_budget=10000,
             window_size=window_size,
-            use_alpha=use_alpha)
-        
+            use_alpha=use_alpha,
+        )
+
     def receive_permutation(
-        self, result: Result, permutation: str, rank_start: int, rank_end: int, logging: bool = False
+        self,
+        result: Result,
+        permutation: str,
+        rank_start: int,
+        rank_end: int,
+        logging: bool = False,
     ) -> Result:
         """
         Processes and applies a permutation to the ranking results.
@@ -54,13 +62,14 @@ def receive_permutation(
             Result: The updated result object with the new ranking order applied.
 
         Note:
-            This function assumes that the permutation string has reasoning generated by Rank-K-32B preceding 
+            This function assumes that the permutation string has reasoning generated by Rank-K-32B preceding
             the sequence of integers separated by spaces.
-            The function would take the last line of input, and call receive_permutation function from 
-            the superclass. 
+            The function would take the last line of input, and call receive_permutation function from
+            the superclass.
         """
 
         # Remove all the reasoning, and take only the order
         permutation = permutation.strip().split("\n")[-1]
-        return super().receive_permutation(result, permutation, rank_start, rank_end, logging)
-    
+        return super().receive_permutation(
+            result, permutation, rank_start, rank_end, logging
+        )

src/rank_llm/rerank/reranker.py

@@ -9,7 +9,12 @@
     get_genai_api_key,
     get_openai_api_key,
 )
-from rank_llm.rerank.listwise import RankListwiseOSLLM, SafeGenai, SafeOpenai, RankKReranker
+from rank_llm.rerank.listwise import (
+    RankKReranker,
+    RankListwiseOSLLM,
+    SafeGenai,
+    SafeOpenai,
+)
 from rank_llm.rerank.listwise.rank_fid import RankFiDDistill, RankFiDScore
 from rank_llm.rerank.pairwise.duot5 import DuoT5
 from rank_llm.rerank.pointwise.monot5 import MonoT5
@@ -523,7 +528,7 @@ def create_model_coordinator(
                 model=(model_path),
                 context_size=context_size,
                 prompt_mode=prompt_mode,
-                prompt_template_path = prompt_template_path,
+                prompt_template_path=prompt_template_path,
                 num_few_shot_examples=num_few_shot_examples,
                 device=device,
                 num_gpus=num_gpus,