Merge pull request #7 from XKTZ/main

merge

README.md

@@ -10,7 +10,7 @@
 We offer a suite of rerankers - pointwise models like monoT5 and listwise models with a focus on open source LLMs compatible with [FastChat](https://github.com/lm-sys/FastChat?tab=readme-ov-file#supported-models) (e.g., Vicuna, Zephyr, etc.) or [vLLM](https://https://github.com/vllm-project/vllm). We also support RankGPT variants, which are proprietary listwise rerankers. Some of the code in this repository is borrowed from [RankGPT](https://github.com/sunnweiwei/RankGPT), [PyGaggle](https://github.com/castorini/pygaggle), and [LiT5](https://github.com/castorini/LiT5)!
 
 # Releases
-current_version = 0.20.1
+current_version = 0.20.2
 
 ## 📟 Instructions
 

pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "rank-llm"
-version = "0.20.1"
+version = "0.20.2"
 description = "A Package for running prompt decoders like RankVicuna"
 readme = "README.md"
 authors = [
@@ -35,7 +35,7 @@ vllm = [
 Homepage = "https://github.com/castorini/rank_llm"
 
 [tool.bumpver]
-current_version = "0.20.1"
+current_version = "0.20.2"
 version_pattern = "MAJOR.MINOR.PATCH"
 commit_message = "Bump version {old_version} -> {new_version}"
 commit = true

src/rank_llm/__init__.py

@@ -1,3 +0,0 @@
-from .retrieve_and_rerank import retrieve_and_rerank
-
-__all__ = ["retrieve_and_rerank"]

src/rank_llm/api/server.py

@@ -3,10 +3,10 @@
 import torch
 from flask import Flask, jsonify, request
 
-from rank_llm import retrieve_and_rerank
 from rank_llm.rerank import PromptMode, get_azure_openai_args, get_openai_api_key
 from rank_llm.rerank.listwise import RankListwiseOSLLM, SafeOpenai
 from rank_llm.retrieve import RetrievalMethod, RetrievalMode
+from rank_llm.retrieve_and_rerank import retrieve_and_rerank
 
 """ API URL FORMAT
 

src/rank_llm/rerank/api_keys.py

@@ -3,10 +3,20 @@
 
 from dotenv import load_dotenv
 
+# Common OpenAI API key paths
+paths = [
+    "OPENAI_API_KEY",
+    "OPEN_AI_API_KEY",
+]
+
 
 def get_openai_api_key() -> str:
     load_dotenv(dotenv_path=f".env.local")
-    return os.getenv("OPEN_AI_API_KEY")
+
+    for path in paths:
+        if os.getenv(path) is not None:
+            return os.getenv(path)
+    return None
 
 
 def get_azure_openai_args() -> Dict[str, str]:

src/rank_llm/rerank/reranker.py

@@ -383,7 +383,9 @@ def extract_kwargs(
     """
     extracted_kwargs = []
     for key, default in keys_and_defaults:
-        value = kwargs.get(key, default)
+        value = kwargs.get(key, None)
+        if value is None:
+            value = default
         if (
             value is not None
             and default is not None

src/rank_llm/retrieve_and_rerank.py

@@ -2,7 +2,6 @@
 from typing import Any, Dict, List, Union
 
 from rank_llm.data import Query, Request
-from rank_llm.evaluation.trec_eval import EvalFunction
 from rank_llm.rerank import IdentityReranker, RankLLM, Reranker
 from rank_llm.rerank.reranker import extract_kwargs
 from rank_llm.retrieve import (
@@ -104,6 +103,8 @@ def retrieve_and_rerank(
             and dataset not in ["dl22", "dl22-passage", "news"]
             and TOPICS[dataset] not in ["dl22", "dl22-passage", "news"]
         ):
+            from rank_llm.evaluation.trec_eval import EvalFunction
+
             print("Evaluating:")
             EvalFunction.eval(["-c", "-m", "ndcg_cut.1", TOPICS[dataset], file_name])
             EvalFunction.eval(["-c", "-m", "ndcg_cut.5", TOPICS[dataset], file_name])

src/rank_llm/scripts/run_rank_llm.py

@@ -9,9 +9,9 @@
 parent = os.path.dirname(parent)
 sys.path.append(parent)
 
-from rank_llm import retrieve_and_rerank
 from rank_llm.rerank import PromptMode
 from rank_llm.retrieve import TOPICS, RetrievalMethod, RetrievalMode
+from rank_llm.retrieve_and_rerank import retrieve_and_rerank
 
 
 def main(args):

test/retrieve/test_ServiceRetriever.py

@@ -1,8 +1,8 @@
 import unittest
 
-from rank_llm import retrieve_and_rerank
 from rank_llm.data import Candidate, Query, Request
 from rank_llm.retrieve import RetrievalMethod, RetrievalMode, ServiceRetriever
+from rank_llm.retrieve_and_rerank import retrieve_and_rerank
 
 
 class TestServiceRetriever(unittest.TestCase):

test/test_retrieve_and_rerank.py

@@ -1,7 +1,7 @@
 import unittest
 from unittest.mock import MagicMock, patch
 
-from rank_llm import retrieve_and_rerank
+from rank_llm.retrieve_and_rerank import retrieve_and_rerank
 
 
 # Anserini API must be hosted at 8081