From d1f94c284681d4cb3db24b529f8bf350c137a7a5 Mon Sep 17 00:00:00 2001
From: hzg0601 <hzg0601@163.com>
Date: Wed, 6 Dec 2023 20:42:27 +0800
Subject: [PATCH] temporarily add reranker

---
 server/reranker/reranker.py | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)
 create mode 100644 server/reranker/reranker.py

diff --git a/server/reranker/reranker.py b/server/reranker/reranker.py
new file mode 100644
index 0000000..5bc702d
--- /dev/null
+++ b/server/reranker/reranker.py
@@ -0,0 +1,19 @@
+from langchain.retrievers.document_compressors import CohereRerank
+from llama_index.postprocessor import SentenceTransformerRerank
+from sentence_transformers import SentenceTransformer,CrossEncoder
+
+model_path = "/root/autodl-tmp/models/bge-reranker-large/"
+instruction = "为这个句子生成表示以用于检索相关文章："
+reranker = SentenceTransformerRerank(
+    top_n=5,
+    model="local:"+model_path,
+)
+
+reranker_model = SentenceTransformer(model_name_or_path=model_path,device="cuda")
+
+reranker_ce = CrossEncoder(model_name=model_path,device="cuda",max_length=1024)
+
+reranker_ce.predict([[],[]])
+
+print("Load reranker")
+