Skip to content

Commit ab1dd9e

Browse files
committed
feat: zs reasoning reranker
feat: working gpt-4.1 script feat: sota reranking with qwen3 feat: sota reranking with qwen3 fix
1 parent 2b53121 commit ab1dd9e

3 files changed

Lines changed: 606 additions & 31 deletions

File tree

openai_ranker.py

Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
from typing import List, Tuple
2+
from together_ranker import TogetherListwiseReranker, TASK2PROMPT
3+
import logging
4+
import os
5+
6+
class OpenAIReranker(TogetherListwiseReranker):
7+
"""Reranker using OpenAI's models, with a prompt optimized for chain-of-thought reasoning."""
8+
9+
def _format_prompt(self, query: str, docs_in_window: List[Tuple[str, str]]) -> str:
10+
"""Formats the user prompt for the LLM with chain-of-thought instructions."""
11+
passages_str = ""
12+
for idx, (doc_id, doc_text) in enumerate(docs_in_window):
13+
# Using 1-based indexing for passages as it's common in prompts
14+
# Not stripping newlines from doc_text to preserve original formatting for the model
15+
passages_str += f"[{idx + 1}] (ID: {doc_id}) {doc_text}\n\n"
16+
17+
# Ensure there's no trailing newline if passages_str is empty, though unlikely.
18+
passages_str = passages_str.strip()
19+
if "gpt-4.1" in self.model_name:
20+
prompt = f"""
21+
Given a query and a list of passages, your task is to re-rank these passages based on their relevance to the query. {TASK2PROMPT[self.task]}
22+
23+
Please perform the following steps:
24+
1. **Understand the Query**: First, carefully read and understand the user's query to identify the core information need.
25+
2. **Analyze Each Passage**: For each passage, critically evaluate its content and determine how well it addresses the query. Consider factors like:
26+
- Directness of the answer
27+
- Completeness of the information
28+
- Presence of supporting evidence or details
29+
- Absence of irrelevant or distracting information
30+
3. **Compare and Contrast**: Compare the passages against each other. Identify which passages are more relevant and why. Note any subtle differences in relevance.
31+
4. **Reasoning for Ranking**: Explicitly state your reasoning for the rank you assign to each passage. Explain why a passage is ranked higher or lower than others. This step-by-step thought process is crucial.
32+
5. **Assign Ranks**: Based on your analysis and reasoning, assign a unique rank to each passage, starting from 1 for the most relevant.
33+
34+
**Output Format:**
35+
Your final output must be a list of ranks, corresponding to the original order of the passages. For example, if there are 3 passages, and you decide the second passage is most relevant, the first is second most relevant, and the third is least relevant, your output should be:
36+
[2] > [1] > [3]
37+
38+
No other text or explanation should be present in the final output, only the list of ranks.
39+
40+
**Query:**
41+
{query}
42+
43+
**Passages:**
44+
{passages_str}
45+
46+
**Your Step-by-Step Reasoning (before the final output list):**
47+
[Provide your detailed thought process here, analyzing each passage and justifying your ranking decisions. This section will not be part of the final output but helps in arriving at the correct ranking.]
48+
49+
**Ranks (only this list will be parsed):**
50+
"""
51+
else:
52+
prompt = f"""
53+
Given a query and a list of passages, your task is to re-rank these passages based on their relevance to the query. {TASK2PROMPT[self.task]}
54+
55+
Please perform the following steps:
56+
1. **Understand the Query**: First, carefully read and understand the user's query to identify the core information need.
57+
2. **Analyze Each Passage**: For each passage, critically evaluate its content and determine how well it addresses the query. Consider factors like:
58+
- Directness of the answer
59+
- Completeness of the information
60+
- Presence of supporting evidence or details
61+
- Absence of irrelevant or distracting information
62+
3. **Compare and Contrast**: Compare the passages against each other. Identify which passages are more relevant and why. Note any subtle differences in relevance.
63+
4. **Assign Ranks**: Based on your analysis and reasoning, assign a unique rank to each passage, starting from 1 for the most relevant.
64+
65+
**Output Format:**
66+
Your final output must be a list of ranks, corresponding to the original order of the passages. For example, if there are 3 passages, and you decide the second passage is most relevant, the first is second most relevant, and the third is least relevant, your output should be:
67+
[2] > [1] > [3]
68+
69+
No other text or explanation should be present in the final output, only the list of ranks.
70+
71+
**Query:**
72+
{query}
73+
74+
**Passages:**
75+
{passages_str}
76+
"""
77+
return prompt
78+
79+
# Example usage (optional, for testing)
80+
if __name__ == '__main__':
81+
logging.basicConfig(level=logging.DEBUG) # Changed to INFO for less verbose default
82+
# Ensure OPENAI_API_KEY is set as an environment variable for testing
83+
# export OPENAI_API_KEY='your_openai_api_key'
84+
85+
# Example documents (doc_id, doc_text)
86+
docs_to_rerank = [
87+
("doc1", "The quick brown fox jumps over the lazy dog."),
88+
("doc2", "A lazy dog sits under the tree."),
89+
("doc3", "Foxes are omnivorous mammals belonging to several genera of the family Canidae."),
90+
("doc4", "The study of dogs is known as cynology."),
91+
("doc5", "Quick feet help the fox escape predators."),
92+
("doc6", "Brown bears are found in North America."),
93+
("doc7", "Lazy rivers flow slowly."),
94+
("doc8", "Jumping requires strong leg muscles."),
95+
("doc9", "The quicksand trapped the unlucky traveler."),
96+
("doc10", "Dog breeds vary greatly in size and temperament."),
97+
]
98+
query_str = "information about quick brown foxes and lazy dogs"
99+
100+
try:
101+
# Retrieve API key from environment variable
102+
# The TogetherListwiseReranker base class should handle the api_key and api_base arguments
103+
# if they are passed to its __init__ method.
104+
# Since OpenAIReranker now uses parent's __init__, we can pass these.
105+
openai_api_key = os.getenv("OPENAI_API_KEY")
106+
if not openai_api_key:
107+
logger = logging.getLogger(__name__)
108+
logger.warning("OPENAI_API_KEY environment variable not set. API calls might fail.")
109+
# Or raise ValueError("OPENAI_API_KEY must be set for testing")
110+
111+
# Instantiate OpenAIReranker
112+
# model_name defaults to "gpt-4-turbo-preview" in parent if not specified.
113+
# The parent __init__ takes api_key and api_base (which is called base_url there)
114+
ranker = OpenAIReranker(
115+
model_name="o3-mini",
116+
window_size=4,
117+
stride=2,
118+
api_key=openai_api_key,
119+
base_url=None,
120+
task="biology" # Provide a task value as parent expects it
121+
)
122+
123+
logger = logging.getLogger(__name__)
124+
logger.info(f"Testing OpenAIReranker with model: {ranker.model_name}")
125+
reranked_ids = ranker.rerank(docs=docs_to_rerank, query=query_str, topk=10)
126+
print("\nReranked Document IDs:", reranked_ids)
127+
128+
# Print reranked docs for verification
129+
reranked_docs_map = {d[0]: d[1] for d in docs_to_rerank}
130+
print("\nReranked Documents:")
131+
for i, doc_id in enumerate(reranked_ids):
132+
print(f"{i+1}. [{doc_id}] {reranked_docs_map[doc_id]}")
133+
134+
except ValueError as e:
135+
print(f"Configuration or Value Error: {e}")
136+
except Exception as e:
137+
# Attempt to get logger from the ranker instance if available, else use global logger
138+
logger_instance = getattr(ranker, 'logger', logging.getLogger(__name__))
139+
logger_instance.error(f"An unexpected error occurred during testing: {e}", exc_info=True)
140+
print(f"An unexpected error occurred: {e}")

rerank.py

Lines changed: 97 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -8,12 +8,16 @@
88
from datasets import load_dataset
99
import torch
1010
from sentence_transformers import CrossEncoder
11+
from concurrent.futures import ThreadPoolExecutor, as_completed
1112

13+
from openai_ranker import OpenAIReranker
1214
from retrievers import calculate_retrieval_metrics
15+
from together_ranker import TogetherListwiseReranker
16+
1317
import functools
1418
import logging
1519
logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
16-
datefmt='%m/%d/%Y %H:%M:%S')
20+
datefmt='%m/%d/%Y %H:%M:%S', level=logging.INFO)
1721
logger = logging.getLogger(__name__)
1822
logger.setLevel(logging.INFO)
1923

@@ -197,18 +201,52 @@ def rerank(self, docs, query, topk):
197201
return ranking
198202

199203

204+
def rerank_single_query(qid, scores, model, documents, examples, args):
205+
logger.debug(f"Reranking qid: {qid}")
206+
try:
207+
sorted_scores = sorted(scores.items(), key=lambda x: x[1], reverse=True)[:args.input_k]
208+
209+
if isinstance(model, TogetherListwiseReranker):
210+
# Together ranker expects List[Tuple[str, str]]
211+
docs_for_rerank = [(did, documents[did]) for did, _ in sorted_scores]
212+
elif isinstance(model, (ClaudeModel, OpenAIModel)):
213+
# Other models might expect List[List[str, str]] - adjust if needed
214+
docs_for_rerank = [[did, documents[did]] for did, _ in sorted_scores]
215+
else:
216+
# Default or handle other model types if necessary
217+
docs_for_rerank = [(did, documents[did]) for did, _ in sorted_scores]
218+
219+
reranked_ids = model.rerank(docs=docs_for_rerank, query=examples[qid]['query'], topk=args.k)
220+
221+
# Assign descending scores based on the new order
222+
final_scores = {doc_id: args.k - i for i, doc_id in enumerate(reranked_ids)}
223+
logger.debug(f"Finished reranking qid: {qid}")
224+
return qid, final_scores
225+
except Exception as e:
226+
logger.error(f"Error reranking qid {qid}: {e}", exc_info=True)
227+
# Return qid and empty scores on error to avoid blocking others
228+
return qid, {}
229+
230+
200231
if __name__=='__main__':
201232
parser = argparse.ArgumentParser()
202233
parser.add_argument('--task', type=str, required=True,
203234
choices=['biology','earth_science','economics','pony','psychology','robotics','theoremqa_questions', "theoremqa_theorems",
204235
'stackoverflow','sustainable_living','aops','leetcode'])
205236
parser.add_argument('--long_context', action='store_true')
206-
parser.add_argument('--llm', type=str, default=None)
237+
parser.add_argument('--llm', type=str, default=None, help="Model name for Claude, OpenAI, or Sentence Transformer rerankers (e.g., 'claude-3-opus-20240229', 'gpt-4-turbo', 'mixedbread-ai/mxbai-rerank-xsmall-v1')")
238+
parser.add_argument('--together_model', type=str, default=None, help="Model name for TogetherListwiseReranker (e.g., 'mistralai/Mixtral-8x7B-Instruct-v0.1')")
239+
parser.add_argument('--openai', action="store_true")
240+
parser.add_argument('--window_size', type=int, default=10, help="Window size for TogetherListwiseReranker sliding window.")
241+
parser.add_argument('--stride', type=int, default=5, help="Stride for TogetherListwiseReranker sliding window.")
207242
parser.add_argument('--score_file', type=str, default=None)
208243
parser.add_argument('--rerank_score_file', type=str, default=None)
209244
parser.add_argument('--input_k', type=int)
210245
parser.add_argument('--k', type=int)
246+
parser.add_argument('--together_api', action="store_true")
247+
parser.add_argument("--workers", type=int, default=6)
211248
args = parser.parse_args()
249+
print(f"Running reranking for {args.task=}")
212250

213251
if os.path.exists(args.rerank_score_file):
214252
print(f"Rerank score file {args.rerank_score_file} already exists.")
@@ -229,38 +267,66 @@ def rerank(self, docs, query, topk):
229267
all_scores = json.load(f)
230268
new_scores = copy.deepcopy(all_scores)
231269

232-
if 'claude' in args.llm:
233-
model = ClaudeModel(version=args.llm)
234-
elif "gpt" in args.llm:
235-
model = OpenAIModel(model_name=args.llm)
236-
else:
237-
model = STReranker(model_name=args.llm)
238-
239-
for qid,scores in tqdm(all_scores.items()):
240-
docs = []
241-
sorted_scores = sorted(scores.items(),key=lambda x:x[1],reverse=True)[:args.input_k]
242-
for did, _ in sorted_scores:
243-
docs.append([did, documents[did]])
244-
245-
if 'claude' in args.llm or "gpt" in args.llm:
246-
new_rank = model.rerank(docs=docs, query=examples[qid]['query'], topk=args.k)
247-
cur_score = {}
248-
if new_rank is None:
249-
# use the original ranks if fail
250-
for rank_id, (did, _) in enumerate(sorted_scores):
251-
cur_score[did] = args.k - rank_id
252-
else:
253-
for rank_id, r in enumerate(new_rank):
254-
cur_score[r] = args.k - rank_id
255-
new_scores[qid] = cur_score
270+
model = None
271+
if args.together_model:
272+
logger.info(f"Using TogetherListwiseReranker with model: {args.together_model}")
273+
if args.openai:
274+
try:
275+
model = OpenAIReranker(model_name=args.together_model,
276+
task=args.task,
277+
window_size=args.window_size,
278+
stride=args.stride,
279+
api_key=os.getenv("OPENAI_API_KEY"),
280+
base_url=None,
281+
together_api=args.together_api
282+
)
283+
except ValueError as e:
284+
logger.error(f"Error initializing OpenAIReranker: {e}")
285+
exit(1)
286+
else:
287+
try:
288+
model = TogetherListwiseReranker(model_name=args.together_model,
289+
task=args.task,
290+
window_size=args.window_size,
291+
stride=args.stride,
292+
together_api=args.together_api
293+
)
294+
except ValueError as e:
295+
logger.error(f"Error initializing TogetherListwiseReranker: {e}")
296+
exit(1)
297+
elif args.llm:
298+
if 'claude' in args.llm:
299+
logger.info(f"Using ClaudeModel with version: {args.llm}")
300+
model = ClaudeModel(version=args.llm)
301+
elif "gpt" in args.llm:
302+
logger.info(f"Using OpenAIModel with model name: {args.llm}")
303+
model = OpenAIModel(model_name=args.llm)
256304
else:
257-
ctxs = [{'id': did, 'text': documents[did]} for did, _ in sorted_scores]
258-
cur_score = model.rerank(query=examples[qid]['query'], docs=ctxs, topk=args.k)
259-
new_scores[qid] = cur_score
305+
logger.info(f"Using STReranker with model name: {args.llm}")
306+
model = STReranker(model_name=args.llm)
307+
else:
308+
logger.error("No reranker specified. Please provide --llm or --together_model.")
309+
exit(1)
310+
311+
reranked_scores = {}
312+
futures = []
313+
with ThreadPoolExecutor(max_workers=args.workers) as executor:
314+
logger.info(f"Submitting {len(all_scores)} queries for parallel reranking...")
315+
for qid, scores in all_scores.items():
316+
futures.append(executor.submit(rerank_single_query, qid, scores, model, documents, examples, args))
317+
318+
logger.info(f"Waiting for reranking tasks to complete...")
319+
for future in tqdm(as_completed(futures), total=len(futures), desc="Reranking Queries"):
320+
try:
321+
qid_result, final_scores_result = future.result()
322+
if final_scores_result:
323+
reranked_scores[qid_result] = final_scores_result
324+
except Exception as e:
325+
logger.error(f"Error retrieving result from future: {e}", exc_info=True)
260326

261327
os.makedirs(os.path.dirname(args.rerank_score_file), exist_ok=True)
262328
with open(args.rerank_score_file, 'w') as f:
263-
json.dump(new_scores, f, indent=2)
329+
json.dump(reranked_scores, f, indent=2)
264330

265331
if args.long_context:
266332
key = 'gold_ids_long'
@@ -275,6 +341,6 @@ def rerank(self, docs, query, topk):
275341
if i in documents:
276342
ground_truth[e['id']][i] = 0
277343

278-
results = calculate_retrieval_metrics(results=new_scores, qrels=ground_truth)
344+
results = calculate_retrieval_metrics(results=reranked_scores, qrels=ground_truth)
279345
with open(args.rerank_score_file.replace(".json", "_results.json"), 'w') as f:
280346
json.dump(results, f, indent=2)

0 commit comments

Comments
 (0)