From 99d28e29f9b8fab9dfc38271b01f2749894c9ce5 Mon Sep 17 00:00:00 2001 From: Robert Washbourne Date: Wed, 2 Oct 2024 19:58:43 -0400 Subject: [PATCH] allow enabling rerankers raw_scores (#393) --- libs/infinity_emb/infinity_emb/fastapi_schemas/pymodels.py | 1 + libs/infinity_emb/infinity_emb/infinity_server.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/libs/infinity_emb/infinity_emb/fastapi_schemas/pymodels.py b/libs/infinity_emb/infinity_emb/fastapi_schemas/pymodels.py index 53e506ab..2541f44b 100644 --- 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 index a3a4a9b0..512d3096 100644 --- 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