Repository: superlinear-ai/raglite (language: python)

Issue:
Remove patch and upgrade rerankers when it fixes incorrect FlashRank reranking
When rerankers fixes https://github.com/AnswerDotAI/rerankers/issues/39, remove the current patch[^1] and upgrade it.

[^1]: https://github.com/superlinear-ai/raglite/blob/main/src/raglite/_flashrank.py


Discussion / hints:

Fix patch:
diff --git a/poetry.lock b/poetry.lock
index 31776eb..4cbc267 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -4969,13 +4969,13 @@ use-chardet-on-py3 = ["chardet (>=3.0.2,<6)"]
 
 [[package]]
 name = "rerankers"
-version = "0.5.3"
+version = "0.6.0"
 description = "A unified API for various document re-ranking models."
 optional = false
 python-versions = ">=3.8"
 files = [
-    {file = "rerankers-0.5.3-py3-none-any.whl", hash = "sha256:773aef2f7db371ec9c7b38b6652a78b7fdb96bb7999dc12817e882104c3e24e0"},
-    {file = "rerankers-0.5.3.tar.gz", hash = "sha256:415bae5e9aaa71e1c2e9932839cc84ff8d939afcf31833bdb398c4bb3cbc67a3"},
+    {file = "rerankers-0.6.0-py3-none-any.whl", hash = "sha256:2d70b4a20b040e87b94c46cc1dff817ddf86488949e477a2508eec08d79976f7"},
+    {file = "rerankers-0.6.0.tar.gz", hash = "sha256:3073456ecc6c704159b39138f30071c8cca15588315bfb3759ea8112ad5bcf1d"},
 ]
 
 [package.dependencies]
@@ -4984,14 +4984,15 @@ pydantic = "*"
 tqdm = "*"
 
 [package.extras]
-all = ["flash-attn", "flashrank", "litellm", "nmslib-metabrainz", "protobuf", "rank-llm", "requests", "sentencepiece", "torch", "transformers"]
+all = ["accelerate (>=0.26.0)", "flash-attn", "flashrank", "litellm", "nmslib-metabrainz", "peft (>=0.13.0)", "pillow", "protobuf", "rank-llm", "requests", "sentencepiece", "torch", "transformers (>=4.45.0)"]
 api = ["requests"]
 dev = ["ipyprogress", "ipython", "ir-datasets", "isort", "pytest", "ranx", "ruff", "srsly"]
 flashrank = ["flashrank"]
 gpt = ["litellm"]
-llmlayerwise = ["flash-attn", "protobuf", "sentencepiece", "torch", "transformers"]
+llmlayerwise = ["flash-attn", "protobuf", "sentencepiece", "torch", "transformers (>=4.45.0)"]
+monovlm = ["accelerate (>=0.26.0)", "flash-attn", "peft (>=0.13.0)", "pillow", "protobuf", "sentencepiece", "torch", "transformers (>=4.45.0)"]
 rankllm = ["nmslib-metabrainz", "rank-llm"]
-transformers = ["protobuf", "sentencepiece", "torch", "transformers"]
+transformers = ["protobuf", "sentencepiece", "torch", "transformers (>=4.45.0)"]
 
 [[package]]
 name = "rich"
@@ -6812,4 +6813,4 @@ ragas = ["ragas"]
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.10,<4.0"
-content-hash = "0b68b835b8d0a256c99deb2e0497f912799be1ce022bab73414bcd9d74bfb7ec"
+content-hash = "bd140a94bea25c626ad990faf55886d0ca589a9524125a0e7d1aa607e9aa1609"
diff --git a/pyproject.toml b/pyproject.toml
index 824ff92..64b8765 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -39,7 +39,7 @@ pydantic = ">=2.7.0"
 pynndescent = ">=0.5.12"
 # Reranking:
 langdetect = ">=1.0.9"
-rerankers = { extras = ["flashrank"], version = ">=0.5.3" }
+rerankers = { extras = ["flashrank"], version = ">=0.6.0" }
 # Storage:
 pg8000 = ">=1.31.2"
 sqlmodel-slim = ">=0.0.18"
diff --git a/src/raglite/_config.py b/src/raglite/_config.py
index c3e8fd6..9c73204 100644
--- a/src/raglite/_config.py
+++ b/src/raglite/_config.py
@@ -8,11 +8,10 @@
 from llama_cpp import llama_supports_gpu_offload
 from sqlalchemy.engine import URL
 
-from raglite._flashrank import PatchedFlashRankRanker as FlashRankRanker
-
 # Suppress rerankers output on import until [1] is fixed.
 # [1] https://github.com/AnswerDotAI/rerankers/issues/36
 with contextlib.redirect_stdout(StringIO()):
+    from rerankers.models.flashrank_ranker import FlashRankRanker
     from rerankers.models.ranker import BaseRanker
 
 
diff --git a/src/raglite/_flashrank.py b/src/raglite/_flashrank.py
deleted file mode 100644
index f69d655..0000000
--- a/src/raglite/_flashrank.py
+++ /dev/null
@@ -1,41 +0,0 @@
-"""Patched version of FlashRankRanker that fixes incorrect reranking [1].
-
-[1] https://github.com/AnswerDotAI/rerankers/issues/39
-"""
-
-import contextlib
-from io import StringIO
-from typing import Any
-
-from flashrank import RerankRequest
-
-# Suppress rerankers output on import until [1] is fixed.
-# [1] https://github.com/AnswerDotAI/rerankers/issues/36
-with contextlib.redirect_stdout(StringIO()):
-    from rerankers.documents import Document
-    from rerankers.models.flashrank_ranker import FlashRankRanker
-    from rerankers.results import RankedResults, Result
-    from rerankers.utils import prep_docs
-
-
-class PatchedFlashRankRanker(FlashRankRanker):
-    def rank(
-        self,
-        query: str,
-        docs: str | list[str] | Document | list[Document],
-        doc_ids: list[str] | list[int] | None = None,
-        metadata: list[dict[str, Any]] | None = None,
-    ) -> RankedResults:
-        docs = prep_docs(docs, doc_ids, metadata)
-        passages = [{"id": doc_idx, "text": doc.text} for doc_idx, doc in enumerate(docs)]
-        rerank_request = RerankRequest(query=query, passages=passages)
-        flashrank_results = self.model.rerank(rerank_request)
-        ranked_results = [
-            Result(
-                document=docs[result["id"]],  # This patches the incorrect ranking in the original.
-                score=result["score"],
-                rank=idx + 1,
-            )
-            for idx, result in enumerate(flashrank_results)
-        ]
-        return RankedResults(results=ranked_results, query=query, has_scores=True)


Test patch:
diff --git a/tests/test_rerank.py b/tests/test_rerank.py
index 7a10ea2..20357c7 100644
--- a/tests/test_rerank.py
+++ b/tests/test_rerank.py
@@ -1,11 +1,11 @@
 """Test RAGLite's reranking functionality."""
 
 import pytest
+from rerankers.models.flashrank_ranker import FlashRankRanker
 from rerankers.models.ranker import BaseRanker
 
 from raglite import RAGLiteConfig, hybrid_search, rerank_chunks, retrieve_chunks
 from raglite._database import Chunk
-from raglite._flashrank import PatchedFlashRankRanker as FlashRankRanker
 
 
 @pytest.fixture(