Create localrerank.py

localrerank.py

@@ -0,0 +1,76 @@
+import os
+import numpy as np
+import logging
+import uvicorn
+import datetime
+from fastapi import FastAPI, Security, HTTPException
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+from FlagEmbedding import FlagReranker
+from pydantic import Field, BaseModel, validator
+from typing import Optional, List
+
+app = FastAPI()
+security = HTTPBearer()
+
+#环境变量传入
+sk_key = os.environ.get('sk-key', 'sk-aaabbbcccdddeeefffggghhhiiijjjkkk')
+
+class QADocs(BaseModel):
+    query: Optional[str]
+    documents: Optional[List[str]]
+
+
+class Singleton(type):
+    def __call__(cls, *args, **kwargs):
+        if not hasattr(cls, '_instance'):
+            cls._instance = super().__call__(*args, **kwargs)
+        return cls._instance
+
+
+RERANK_MODEL_PATH = "Qwen/Qwen3-Reranker-4B"
+
+class ReRanker(metaclass=Singleton):
+    def __init__(self, model_path):
+        self.reranker = FlagReranker(model_path, use_fp16=False)
+
+    def compute_score(self, pairs: List[List[str]]):
+        if len(pairs) > 0:
+            result = self.reranker.compute_score(pairs, normalize=True)
+            if isinstance(result, float):
+                result = [result]
+            return result
+        else:
+            return None
+
+class Chat(object):
+    def __init__(self, rerank_model_path: str = RERANK_MODEL_PATH):
+        self.reranker = ReRanker(rerank_model_path)
+
+    def fit_query_answer_rerank(self, query_docs: QADocs) -> List:
+        if query_docs is None or len(query_docs.documents) == 0:
+            return []
+
+        pair = [[query_docs.query, doc] for doc in query_docs.documents]
+        scores = self.reranker.compute_score(pair)
+
+        new_docs = []
+        for index, score in enumerate(scores):
+            new_docs.append({"index": index, "text": query_docs.documents[index], "score": score})
+        results = [{"index": documents["index"], "relevance_score": documents["score"]} for documents in list(sorted(new_docs, key=lambda x: x["score"], reverse=True))]
+        return results
+
+@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}
+    except Exception as e:
+        print(f"报错：\n{e}")
+        return {"error": "重排出错"}
+
+if __name__ == "__main__":
+    uvicorn.run("localrerank:app", host='0.0.0.0', port=7860, workers=1)