First implementation of the ReRanker endpoint.

Author: kalaspuffar

app/models.py

@@ -42,3 +42,10 @@ class QueryMultipleBody(BaseModel):
     query: str
     file_ids: List[str]
     k: int = 4
+
+
+class QueryMultipleDocs(BaseModel):
+    query: str
+    docs: List[str]
+    config: dict
+    k: int = 4

app/routes/document_routes.py

@@ -6,6 +6,7 @@
 import aiofiles.os
 from shutil import copyfileobj
 from typing import List, Iterable
+from rerankers import Reranker, Document as ReRankDocument
 from fastapi import (
     APIRouter,
     Request,
@@ -29,6 +30,7 @@
     QueryRequestBody,
     DocumentResponse,
     QueryMultipleBody,
+    QueryMultipleDocs,
 )
 from app.services.vector_store.async_pg_vector import AsyncPgVector
 from app.utils.document_loader import (
@@ -702,7 +704,6 @@ async def query_embeddings_by_file_ids(request: Request, body: QueryMultipleBody
         )
         raise HTTPException(status_code=500, detail=str(e))
 
-
 @router.post("/text")
 async def extract_text_from_file(
     request: Request,
@@ -765,3 +766,42 @@ async def extract_text_from_file(
             )
     finally:
         await cleanup_temp_file_async(temp_file_path)
+
+@router.post("/rerank")
+async def rerank_documents_by_query(request: Request, body: QueryMultipleDocs):
+    try:
+        rk = Reranker(
+            body.config.get("model_name", "flashrank"),
+            model_type=body.config.get("model_type"),
+            lang=body.config.get("lang"),
+            api_provider=body.config.get("api_provider"),
+            api_key=body.config.get("api_key"),
+        )
+
+        docs = []
+        for i, d in enumerate(body.docs):
+            if isinstance(d, str):
+                docs.append(ReRankDocument(text=d, doc_id=i))
+            else:
+                docs.append(
+                    ReRankDocument(
+                        text=d.get("text", ""),
+                        doc_id=d.get("doc_id", i),
+                        metadata=d.get("metadata", {}) or {},
+                    )
+                )
+
+        top_k = body.k
+
+        results = rk.rank(query=body.query, docs=docs)
+        items = results.top_k(top_k) if top_k else results
+
+        return [[getattr(r.document, "text", None), r.score] for r in items]
+    except Exception as e:
+        logger.error(
+            "Error in reranking documents | Query: %s | Error: %s | Traceback: %s",
+            body.query,
+            str(e),
+            traceback.format_exc(),
+        )
+        raise HTTPException(status_code=500, detail=str(e))

docker-compose.yaml

@@ -17,8 +17,10 @@ services:
       - DB_PORT=5432
     ports:
       - "8000:8000"
+    runtime: nvidia
     volumes:
       - ./uploads:/app/uploads
+      - ~/.cache/huggingface:/root/.cache/huggingface:rw
     depends_on:
       - db
     env_file:

requirements.txt

@@ -38,3 +38,4 @@ python-pptx==1.0.2
 xlrd==2.0.2
 pydantic==2.9.2
 chardet==5.2.0
+rerankers[transformers]==0.6.0