feat: Add BgeRerank test for document reranking

test/k_store_test.py

@@ -15,6 +15,7 @@ from src.config.consts import (
     INDEX_NAME
 )
 from src.loader.callback import BaseCallback
+from src.server.rerank import BgeRerank
 
 sys.path.append("../")
 
@@ -81,6 +82,15 @@ def test_faiss_load():
     print(vecstore_faiss.join_document(vecstore_faiss.get_text_similarity("我国什么时候全面开放低空领域")))
 
 
+def test_reranker():
+    from langchain_core.documents import Document
+    reranker = BgeRerank(model_name="BAAI/bge-reranker-large")
+    docs = ["低空经济是指在 200 米的空域相关业务","我国什么时候全面开放低空领域","今天早上雨很大"]
+    docs2 = [Document(page_content=d) for d in docs]
+    print(reranker.bge_rerank("我国什么时候全面开放低空领域", docs))
+    print(reranker.compress_documents(docs2, "我国什么时候全面开放低空领域"))
+
 if __name__ == "__main__":
-    test_faiss_from_dir()
-    test_faiss_load()
+    # test_faiss_from_dir()
+    # test_faiss_load()
+    test_reranker()