Minor Modification

Author: Yaohui2019

src/rank_llm/rerank/listwise/rankk_reranker.py

@@ -1,7 +1,6 @@
 from typing import Optional
 
 from rank_llm.data import Result
-from rank_llm.rerank import PromptMode
 from rank_llm.rerank.listwise import RankListwiseOSLLM
 
 
@@ -10,7 +9,6 @@ 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",
@@ -24,14 +22,13 @@ def __init__(
         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_gpus=num_gpus,
             variable_passages=variable_passages,
             is_thinking=True,
-            reasoning_token_budget=10000,
+            reasoning_token_budget=8000,
             window_size=window_size,
             use_alpha=use_alpha,
         )

src/rank_llm/rerank/prompt_templates/rank_k_template.yaml

@@ -1,5 +1,5 @@
 method: "singleturn_listwise"
-prefix: |-
+prefix: |+
   Determine a ranking of the passages based on how relevant they are to the query. 
   If the query is a question, how relevant a passage is depends on how well it answers the question. 
   If not, try analyze the intent of the query and assess how well each passage satisfy the intent. 
@@ -13,5 +13,7 @@ prefix: |-
   Output only the ordering without any other text.
 
   Query: {query}
-body: "\n\n[{rank}] {candidate}"
-suffix: ""
+
+body: "[{rank}] {candidate}\n\n"
+output_validation_regex: r"^\[\d+\]( (>|=) \[\d+\])*$"
+output_extraction_regex: r"\[(\d+)\]"

src/rank_llm/rerank/reranker.py

@@ -474,7 +474,6 @@ def create_model_coordinator(
             print(f"Loading {model_path} ...")
             keys_and_defaults = [
                 ("context_size", 4096),
-                ("prompt_mode", PromptMode.RANK_GPT),
                 (
                     "prompt_template_path",
                     "src/rank_llm/rerank/prompt_templates/rank_k_template.yaml",
@@ -490,7 +489,6 @@ def create_model_coordinator(
             ]
             [
                 context_size,
-                prompt_mode,
                 prompt_template_path,
                 num_few_shot_examples,
                 device,
@@ -505,7 +503,6 @@ def create_model_coordinator(
             model_coordinator = RankKReranker(
                 model=(model_path),
                 context_size=context_size,
-                prompt_mode=prompt_mode,
                 prompt_template_path=prompt_template_path,
                 num_few_shot_examples=num_few_shot_examples,
                 device=device,