Update app.py

app.py

@@ -70,37 +70,33 @@ class RerankerResponse(BaseModel):
 async def rerank_documents(request: RerankerRequest):
     """
     Reranks the given list of documents based on their relevance to the query.
-
     - **query**: The input query string.
     - **documents**: A list of documents to be reranked.
     - **return_documents**: Whether to include document content in results.
     - **top_k**: Number of top-ranked documents to return.
-
     Returns:
     - A list of ranked documents with scores and indexes.
     """
     try:
         # Prepare model input
-        inputs = [[request.query, doc] for doc in request.documents]
-
-        # Get ranking scores
-        scores = model.predict(inputs)
-
-        # Sort scores and get top-k results
-        ranked_indices = sorted(
-            range(len(scores)), key=lambda i: scores[i], reverse=True
-        )[: request.top_k]
-
-        # Format results
+        results = model.rank(
+            request.query,
+            request.documents,
+            return_documents=request.return_documents,
+            top_k=request.top_k,
+        )
+
+        # Format the results based on the model's output
         formatted_results = [
             RankedResult(
-                score=scores[i],
-                index=i,
-                document=request.documents[i] if request.return_documents else None,
+                score=result["score"],
+                index=result["corpus_id"],
+                document=result["text"] if request.return_documents else None,
             )
-            for i in ranked_indices
+            for result in results
         ]
 
+        # Format results
         return RerankerResponse(results=formatted_results)
 
     except Exception as e: