diff --git a/src/core/knowledgebase.py b/src/core/knowledgebase.py index 6d68087e..4b16b58d 100644 --- a/src/core/knowledgebase.py +++ b/src/core/knowledgebase.py @@ -67,7 +67,7 @@ class KnowledgeBase: logger.error(f"Collection {collection_name} not found, create it") # self.add_collection(collection_name) - vectors = self.embed_model.encode(docs) + vectors = self.embed_model.batch_encode(docs) data = [{ "id": int(random.random() * 1e12), diff --git a/src/core/retriever.py b/src/core/retriever.py index ae649fad..b31f7ef8 100644 --- a/src/core/retriever.py +++ b/src/core/retriever.py @@ -63,8 +63,8 @@ class Retriever: def query_classification(self, query): """判断是否需要查询 - - 对于完全基于用户给定信息的任务,称之为“足够”“sufficient”,不需要检索; - - 否则,称之为“不足”“insufficient”,可能需要检索, + - 对于完全基于用户给定信息的任务,称之为"足够""sufficient",不需要检索; + - 否则,称之为"不足""insufficient",可能需要检索, """ raise NotImplementedError diff --git a/src/models/embedding.py b/src/models/embedding.py index 136b130b..853a9020 100644 --- a/src/models/embedding.py +++ b/src/models/embedding.py @@ -28,6 +28,7 @@ class RemoteEmbeddingModel: embed_state = {} def batch_encode(self, messages, batch_size=20): + logger.info(f"Batch encoding {len(messages)} messages") data = [] if len(messages) > batch_size: diff --git a/src/routers/data_router.py b/src/routers/data_router.py index ca59e9fa..6a2bf2ce 100644 --- a/src/routers/data_router.py +++ b/src/routers/data_router.py @@ -99,7 +99,6 @@ async def get_graph_info(): @data.get("/graph/node") async def get_graph_node(entity_name: str): - logger.debug(f"Get graph node {entity_name}") result = startup.dbm.graph_base.query_node(entity_name=entity_name) return {"result": startup.retriever.format_query_results(result), "message": "success"}