Merge main

BlueBrain · Sep 11, 2024 · 3db6019 · 3db6019
2 parents d7449b0 + 5fa7dd3
commit 3db6019
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 6 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 - Use OpenAI response_format instead of separators in the prompt.
+- Switch to cohere reranker v3 and `retriever_k = 500`.
 
 ## [v0.0.5]
 

diff --git a/pyproject.toml b/pyproject.toml
@@ -17,7 +17,7 @@ dependencies = [
     "cohere",
     "elasticsearch >= 8.5",
     "elasticsearch-dsl",
-    "fastapi",
+    "fastapi <= 0.112.0",
     "fastapi-pagination",
     "httpx",
     "python-dotenv",

diff --git a/src/scholarag/app/schemas.py b/src/scholarag/app/schemas.py
@@ -50,7 +50,7 @@ class PassthroughRequest(BaseModel):
 class RetrievalRequest(PassthroughRequest, extra="forbid"):
     """Request for the raw retrieval endpoint."""
 
-    retriever_k: int = Field(700, ge=1, le=2000)
+    retriever_k: int = Field(500, ge=1, le=1000)
     use_reranker: bool = True
     reranker_k: int = Field(8, ge=1, le=100)
 

diff --git a/src/scholarag/services/cohere_reranker.py b/src/scholarag/services/cohere_reranker.py
@@ -51,12 +51,12 @@ def run(
             List of dictionaries containing in score order the text and the associated score.
         """
         response = self.client.rerank(
-            model="rerank-english-v2.0",
+            model="rerank-english-v3.0",
             query=query,
             documents=contexts,
             return_documents=True,
             top_n=len(contexts),
-            max_chunks_per_doc=10_000 // len(contexts),
+            max_chunks_per_doc=1000 // len(contexts),
         )
         reranked_contexts = [
             {
@@ -88,12 +88,12 @@ async def arun(
             List of dictionaries containing in score order the text and the associated score.
         """
         response = await self.async_client.rerank(
-            model="rerank-english-v2.0",
+            model="rerank-english-v3.0",
             query=query,
             documents=contexts,
             return_documents=True,
             top_n=len(contexts),
-            max_chunks_per_doc=10_000 // len(contexts),
+            max_chunks_per_doc=1000 // len(contexts),
         )
 
         reranked_contexts = [