allow enabling rerankers raw_scores (#393)

michaelfeil · Oct 2, 2024 · 99d28e2 · 99d28e2
1 parent bb39cbd
commit 99d28e2
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/libs/infinity_emb/infinity_emb/fastapi_schemas/pymodels.py b/libs/infinity_emb/infinity_emb/fastapi_schemas/pymodels.py
@@ -183,6 +183,7 @@ class RerankInput(BaseModel):
         **ITEMS_LIMIT,
     )
     return_documents: bool = False
+    raw_scores: bool = False
     model: str = "default/not-specified"
 
 

diff --git a/libs/infinity_emb/infinity_emb/infinity_server.py b/libs/infinity_emb/infinity_emb/infinity_server.py
@@ -290,7 +290,7 @@ async def _rerank(data: RerankInput):
             start = time.perf_counter()
 
             scores, usage = await engine.rerank(
-                query=data.query, docs=data.documents, raw_scores=False
+                query=data.query, docs=data.documents, raw_scores=data.raw_scores
             )
 
             duration = (time.perf_counter() - start) * 1000