Update apis/chat_api.py

apis/chat_api.py

@@ -26,7 +26,7 @@ from networks.huggingface_streamer import HuggingfaceStreamer
 from networks.huggingchat_streamer import HuggingchatStreamer
 from networks.openai_streamer import OpenaiStreamer
 
-from sentence_transformers import SentenceTransformer
+from sentence_transformers import SentenceTransformer, CrossEncoder
 import tiktoken
 
 class EmbeddingsAPIInference:
@@ -72,6 +72,9 @@ class ChatAPIApp:
             "intfloat/multilingual-e5-large-instruct":EmbeddingsAPIInference("intfloat/multilingual-e5-large-instruct"),
             "mixedbread-ai/mxbai-embed-large-v1":EmbeddingsAPIInference("mixedbread-ai/mxbai-embed-large-v1")
         }
+        self.rerank = {
+            "bge-reranker-v2-m3":CrossEncoder("BAAI/bge-reranker-v2-m3")
+        }
 
     def get_available_models(self):
         return {"object": "list", "data": AVAILABLE_MODELS_DICTS}
@@ -376,6 +379,22 @@ class ChatAPIApp:
             return {"embedding": embeddings}#.tolist()}
         except ValueError as e:
             raise HTTPException(status_code=400, detail=str(e))
+
+
+    class RerankRequest(BaseModel):
+        model: str
+        input: str
+        document: list
+        top_k: int
+        return_documents: bool
+
+    def get_rerank(self, request: RerankRequest, api_key: str = Depends(extract_api_key)):
+        return self.rerank[request.model].rank(
+            request.input, 
+            request.document, 
+            top_k=request.top_k, 
+            return_documents=request.return_documents
+        )
         
 
     def get_readme(self):
@@ -400,6 +419,12 @@ class ChatAPIApp:
                 include_in_schema=include_in_schema,
             )(self.get_available_models)
 
+            self.app.post(
+                prefix+"/rerank",
+                summary="Rerank documents",
+                include_in_schema=include_in_schema,
+            )(self.get_rerank)
+
             self.app.post(
                 prefix + "/chat/completions",
                 summary="OpenAI Chat completions in conversation session",