Commit 1691bed8 by 文靖昊

Merge remote-tracking branch 'origin/geo' into geo

parents 4f95b54c e7755af4
...@@ -15,6 +15,7 @@ from src.config.consts import ( ...@@ -15,6 +15,7 @@ from src.config.consts import (
INDEX_NAME INDEX_NAME
) )
from src.loader.callback import BaseCallback from src.loader.callback import BaseCallback
from src.server.rerank import BgeRerank
sys.path.append("../") sys.path.append("../")
...@@ -81,6 +82,15 @@ def test_faiss_load(): ...@@ -81,6 +82,15 @@ def test_faiss_load():
print(vecstore_faiss.join_document(vecstore_faiss.get_text_similarity("我国什么时候全面开放低空领域"))) 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__": if __name__ == "__main__":
test_faiss_from_dir() # test_faiss_from_dir()
test_faiss_load() # test_faiss_load()
test_reranker()
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment