Update localrerank.py

localrerank.py

@@ -5,7 +5,7 @@ import uvicorn
 import datetime
 from fastapi import FastAPI, Security, HTTPException
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
-from FlagEmbedding import FlagReranker
+from FlagEmbedding import FlagAutoReranker
 from pydantic import Field, BaseModel, validator
 from typing import Optional, List
 
@@ -44,7 +44,11 @@ class ReRanker(metaclass=Singleton):
 
 class Chat(object):
     def __init__(self, rerank_model_path: str = RERANK_MODEL_PATH):
-        self.reranker = ReRanker(rerank_model_path)
+        self.reranker = FlagAutoReranker.from_finetuned(
+            model_path,
+            use_fp16=False,
+            trust_remote_code=True,
+        )
 
     def fit_query_answer_rerank(self, query_docs: QADocs) -> List:
         if query_docs is None or len(query_docs.documents) == 0:
@@ -59,12 +63,13 @@ class Chat(object):
         results = [{"index": documents["index"], "relevance_score": documents["score"]} for documents in list(sorted(new_docs, key=lambda x: x["score"], reverse=True))]
         return results
 
+chat = Chat()
+
 @app.post('/v1/rerank')
 async def handle_post_request(docs: QADocs, credentials: HTTPAuthorizationCredentials = Security(security)):
     token = credentials.credentials
     if sk_key is not None and token != sk_key:
         raise HTTPException(status_code=401, detail="Invalid token")
-    chat = Chat()
     try:
         results = chat.fit_query_answer_rerank(docs)
         return {"results": results}