add top_k parameter to rerank endpoint

libs/infinity_emb/infinity_emb/engine.py

@@ -152,14 +152,21 @@ async def embed(
         return embeddings, usage
 
     async def rerank(
-        self, *, query: str, docs: list[str], raw_scores: bool = False
+        self,
+        *,
+        query: str,
+        docs: list[str],
+        raw_scores: bool = False,
+        top_k: Optional[int] = None,
     ) -> tuple[list[float], int]:
         """rerank multiple sentences
 
         Kwargs:
             query (str): query to be reranked
             docs (list[str]): docs to be reranked
             raw_scores (bool): return raw scores instead of sigmoid
+            top_k (Optional[int]): number of top scores to return after reranking
+                if top_k is None, <= 0 or out of range, all scores are returned
 
         Raises:
             ValueError: raised if engine is not started yet
@@ -172,7 +179,10 @@ async def rerank(
         """
         self._assert_running()
         scores, usage = await self._batch_handler.rerank(
-            query=query, docs=docs, raw_scores=raw_scores
+            query=query,
+            docs=docs,
+            raw_scores=raw_scores,
+            top_k=top_k,
         )
 
         return scores, usage

libs/infinity_emb/infinity_emb/fastapi_schemas/pymodels.py

@@ -257,6 +257,7 @@ class RerankInput(BaseModel):
     return_documents: bool = False
     raw_scores: bool = False
     model: str = "default/not-specified"
+    top_k: Optional[int] = Field(default=None, gt=0)
 
 
 class _ReRankObject(BaseModel):

libs/infinity_emb/infinity_emb/inference/batch_handler.py

@@ -9,7 +9,7 @@
 import time
 from concurrent.futures import ThreadPoolExecutor
 from queue import Queue
-from typing import Any, List, Sequence, Set, Union
+from typing import Any, List, Sequence, Set, Union, Optional
 
 import numpy as np
 
@@ -147,14 +147,20 @@ async def embed(
         return embeddings, usage
 
     async def rerank(
-        self, query: str, docs: list[str], raw_scores: bool = False
+        self,
+        query: str,
+        docs: list[str],
+        raw_scores: bool = False,
+        top_k: Optional[int] = None,
     ) -> tuple[list[float], int]:
         """Schedule a query to be reranked with documents. Awaits until reranked.
 
         Args:
             query (str): query for reranking
             docs (list[str]): documents to be reranked
             raw_scores (bool): return raw scores instead of sigmoid
+            top_k (Optional[int]): number of top scores to return after reranking
+                if top_k is None, <= 0 or out of range, all scores are returned
 
         Raises:
             ModelNotDeployedError: If loaded model does not expose `embed`
@@ -172,6 +178,9 @@ async def rerank(
         rerankables = [ReRankSingle(query=query, document=doc) for doc in docs]
         scores, usage = await self._schedule(rerankables)
 
+        if top_k is not None and top_k > 0:
+            scores = scores[:top_k]
+
         if not raw_scores:
             # perform sigmoid on scores
             scores = (1 / (1 + np.exp(-np.array(scores)))).tolist()

libs/infinity_emb/infinity_emb/infinity_server.py

@@ -392,7 +392,10 @@ async def _rerank(data: RerankInput):
             start = time.perf_counter()
 
             scores, usage = await engine.rerank(
-                query=data.query, docs=data.documents, raw_scores=data.raw_scores
+                query=data.query,
+                docs=data.documents,
+                raw_scores=data.raw_scores,
+                top_k=data.top_k,
             )
 
             duration = (time.perf_counter() - start) * 1000

libs/infinity_emb/tests/end_to_end/test_torch_reranker.py

@@ -1,5 +1,6 @@
 import pytest
 import torch
+import sys
 from asgi_lifespan import LifespanManager
 from httpx import AsyncClient
 from transformers import pipeline  # type: ignore[import-untyped]
@@ -89,6 +90,61 @@ async def test_reranker(client, model_base, helpers):
     for i, pred in enumerate(predictions):
         assert abs(rdata_results[i]["relevance_score"] - pred["score"]) < 0.01
 
+@pytest.mark.anyio
+async def test_reranker_top_k(client):
+    query = "Where is the Eiffel Tower located?"
+    documents = [
+        "The Eiffel Tower is located in Paris, France",
+        "The Eiffel Tower is located in the United States.",
+        "The Eiffel Tower is located in the United Kingdom.",
+    ]
+
+    response = await client.post(
+        f"{PREFIX}/rerank",
+        json={"model": MODEL, "query": query, "documents": documents, "top_k": 1},
+    )
+    assert response.status_code == 200
+    rdata = response.json()
+    rdata_results = rdata["results"]
+    assert len(rdata_results) == 1
+
+    response = await client.post(
+        f"{PREFIX}/rerank",
+        json={"model": MODEL, "query": query, "documents": documents, "top_k": 2},
+    )
+    assert response.status_code == 200
+    rdata = response.json()
+    rdata_results = rdata["results"]
+    assert len(rdata_results) == 2
+
+    response = await client.post(
+        f"{PREFIX}/rerank",
+        json={"model": MODEL, "query": query, "documents": documents, "top_k": sys.maxsize},
+    )
+    assert response.status_code == 200
+    rdata = response.json()
+    rdata_results = rdata["results"]
+    assert len(rdata_results) == len(documents)
+
+@pytest.mark.anyio
+async def test_reranker_invalid_top_k(client):
+    query = "Where is the Eiffel Tower located?"
+    documents = [
+        "The Eiffel Tower is located in Paris, France",
+        "The Eiffel Tower is located in the United States.",
+        "The Eiffel Tower is located in the United Kingdom.",
+    ]
+    response = await client.post(
+        f"{PREFIX}/rerank",
+        json={"model": MODEL, "query": query, "documents": documents, "top_k": -1},
+    )
+    assert response.status_code == 422
+
+    response = await client.post(
+        f"{PREFIX}/rerank",
+        json={"model": MODEL, "query": query, "documents": documents, "top_k": 0},
+    )
+    assert response.status_code == 422
 
 @pytest.mark.anyio
 async def test_reranker_cant_embed_or_classify(client):

libs/infinity_emb/tests/unit_test/test_engine.py

@@ -98,6 +98,50 @@
     np.testing.assert_almost_equal(rankings[:3], [0.83, 0.085, 0.028], decimal=2)
 
 
+@pytest.mark.anyio
+@pytest.mark.parametrize("engine", [InferenceEngine.torch, InferenceEngine.optimum])
+async def test_engine_reranker_top_k(engine):
+    model_unpatched = CrossEncoder(
+        "mixedbread-ai/mxbai-rerank-xsmall-v1",
+    )
+    query = "Where is Paris?"
+    documents = [
+        "Paris is the capital of France.",
+        "Berlin is the capital of Germany.",
+        "You can now purchase my favorite dish",
+    ]
+    engine = AsyncEmbeddingEngine.from_args(
+        EngineArgs(
+            model_name_or_path="mixedbread-ai/mxbai-rerank-xsmall-v1",
+            engine=InferenceEngine.torch,
+            model_warmup=False,
+        )
+    )
+
+    query_docs = [(query, doc) for doc in documents]
+
+    async with engine:
+        rankings, usage = await engine.rerank(query=query, docs=documents, top_k=None)
+    assert len(rankings) == len(documents)
+
+    async with engine:
+        rankings, usage = await engine.rerank(query=query, docs=documents, top_k=-1)
+    assert len(rankings) == len(documents)
+
+    async with engine:
+        rankings, usage = await engine.rerank(query=query, docs=documents, top_k=0)
+    assert len(rankings) == len(documents)
+
+    async with engine:
+        rankings, usage = await engine.rerank(query=query, docs=documents, top_k=3)
+    assert len(rankings) == 3
+
+    async with engine:
+        rankings, usage = await engine.rerank(query=query, docs=documents, top_k=len(documents)+sys.maxsize)
+    assert len(rankings) == len(documents)
+
+
+
 @pytest.mark.anyio
 async def test_async_api_torch_CLASSIFY():
     sentences = ["This is awesome.", "I am depressed."]