新增自定义关键词检索类

backend/app/engine/__init__.py

@@ -11,6 +11,7 @@ from sqlalchemy import create_engine, Engine
 from app.engine.loaders.db import makeDescriptionByEngine
 from app.engine.tools import ToolFactory
 from app.engine.index import get_index
+from app.engine.retriever.CHBM25Retriever import CHBM25Retriever
 from app.settings import get_node_postprocessors
 
 from llama_index.core.retrievers import BaseRetriever
@@ -29,9 +30,6 @@ class HybridRetriever(BaseRetriever):
         filters = None,
         **kwargs: Any,
     ) -> None:
-        from llama_index.retrievers.bm25 import BM25Retriever
-        from nltk.corpus import stopwords
-
         super().__init__(**kwargs)
         self._vector_index = vector_index
         self._embed_model = vector_index._embed_model
@@ -39,9 +37,13 @@ class HybridRetriever(BaseRetriever):
         self._vecRetriever = vector_index.as_retriever(
             similarity_top_k=similarity_top_k,filters = filters
         )
-        self._bm25Retriever = BM25Retriever.from_defaults(similarity_top_k=similarity_top_k,
-                                             nodes=self._vector_index.vector_store.get_nodes(None),
-                                             language=stopwords.words('chinese'))
+
+        STORAGE_DIR = os.getenv("BM_RETRIEVER_PATH", "storage_bm")
+        if os.path.exists(STORAGE_DIR) and len(os.listdir(STORAGE_DIR)) > 0:
+            self._bm25Retriever = CHBM25Retriever.from_persist_dir(STORAGE_DIR)
+        else:
+            bmRetriver = CHBM25Retriever.from_defaults(similarity_top_k=similarity_top_k,nodes=self._vector_index.vector_store.get_nodes(None))
+            bmRetriver.persist(STORAGE_DIR)
         self._alpha = alpha
 
     def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]: