Merge pull request #167 from xerrors/async-for-kg-embedding

优化embedding阶段的并行处理

server/routers/chat_router.py

@@ -21,7 +21,7 @@ async def chat_get():
     return "Chat Get!"
 
 @chat.post("/")
-def chat_post(
+async def chat_post(
         query: str = Body(...),
         meta: dict = Body(None),
         history: list[dict] | None = Body(None),

server/routers/data_router.py

@@ -171,11 +171,15 @@ async def get_graph_nodes(kgdb_name: str, num: int):
 @data.post("/graph/add-by-jsonl")
 async def add_graph_entity(file_path: str = Body(...), kgdb_name: Optional[str] = Body(None)):
     if not config.enable_knowledge_graph:
-        raise HTTPException(status_code=400, detail="Knowledge graph is not enabled")
+        return {"message": "知识图谱未启用", "status": "failed"}
 
     if not file_path.endswith('.jsonl'):
-        raise HTTPException(status_code=400, detail="file_path must be a jsonl file")
+        return {"message": "文件格式错误，请上传jsonl文件", "status": "failed"}
 
-    graph_base.jsonl_file_add_entity(file_path, kgdb_name)
-    return {"message": "Entity successfully added"}
+    try:
+        await graph_base.jsonl_file_add_entity(file_path, kgdb_name)
+        return {"message": "实体添加成功", "status": "success"}
+    except Exception as e:
+        logger.error(f"添加实体失败: {e}, {traceback.format_exc()}")
+        return {"message": f"添加实体失败: {e}", "status": "failed"}
 

src/core/graphbase.py

@@ -106,7 +106,7 @@ class GraphDatabase:
         with self.driver.session() as session:
             session.execute_write(create, triples)
 
-    def txt_add_vector_entity(self, triples, kgdb_name='neo4j'):
+    async def txt_add_vector_entity(self, triples, kgdb_name='neo4j'):
         """添加实体三元组"""
         self.use_database(kgdb_name)
         def _index_exists(tx, index_name):
@@ -140,6 +140,31 @@ class GraphDatabase:
                 }} }};
                 """)
 
+        def _get_nodes_without_embedding(tx, entity_names):
+            """获取没有embedding的节点列表"""
+            # 构建参数字典，将列表转换为"param0"、"param1"等键值对形式
+            params = {f"param{i}": name for i, name in enumerate(entity_names)}
+
+            # 构建查询参数列表
+            param_placeholders = ", ".join([f"${key}" for key in params.keys()])
+
+            # 执行查询
+            result = tx.run(f"""
+            MATCH (n:Entity)
+            WHERE n.name IN [{param_placeholders}] AND n.embedding IS NULL
+            RETURN n.name AS name
+            """, params)
+
+            return [record["name"] for record in result]
+
+        def _batch_set_embeddings(tx, entity_embedding_pairs):
+            """批量设置实体的嵌入向量"""
+            for entity_name, embedding in entity_embedding_pairs:
+                tx.run("""
+                MATCH (e:Entity {name: $name})
+                CALL db.create.setNodeVectorProperty(e, 'embedding', $embedding)
+                """, name=entity_name, embedding=embedding)
+
         # 判断模型名称是否匹配
         cur_embed_info = config.embed_model_names[config.embed_model]
         self.embed_model_name = self.embed_model_name or cur_embed_info.get('name')
@@ -151,18 +176,44 @@ class GraphDatabase:
             session.execute_write(_create_graph, triples)
             logger.info(f"Creating vector index for {kgdb_name} with {config.embed_model}")
             session.execute_write(_create_vector_index, cur_embed_info['dimension'])
-            # NOTE 这里需要异步处理
-            for i, entry in enumerate(triples):
-                logger.debug(f"Adding entity {i+1}/{len(triples)}")
-                embedding_h = self.get_embedding(entry['h'])
-                embedding_t = self.get_embedding(entry['t'])
-                session.execute_write(self.set_embedding, entry['h'], embedding_h)
-                session.execute_write(self.set_embedding, entry['t'], embedding_t)
+
+            # 收集所有需要处理的实体名称，去重
+            all_entities = []
+            for entry in triples:
+                if entry['h'] not in all_entities:
+                    all_entities.append(entry['h'])
+                if entry['t'] not in all_entities:
+                    all_entities.append(entry['t'])
+
+            # 筛选出没有embedding的节点
+            nodes_without_embedding = session.execute_read(_get_nodes_without_embedding, all_entities)
+            if not nodes_without_embedding:
+                logger.info(f"所有实体已有embedding，无需重新计算")
+                return
+
+            logger.info(f"需要为{len(nodes_without_embedding)}/{len(all_entities)}个实体计算embedding")
+
+            # 批量处理实体
+            max_batch_size = 1024  # 限制此部分的主要是内存大小 1024 * 1024 * 4 / 1024 / 1024 = 4GB
+            total_entities = len(nodes_without_embedding)
+
+            for i in range(0, total_entities, max_batch_size):
+                batch_entities = nodes_without_embedding[i:i+max_batch_size]
+                logger.debug(f"Processing entities batch {i//max_batch_size + 1}/{(total_entities-1)//max_batch_size + 1} ({len(batch_entities)} entities)")
+
+                # 批量获取嵌入向量
+                batch_embeddings = await self.aget_embedding(batch_entities)
+
+                # 将实体名称和嵌入向量配对
+                entity_embedding_pairs = list(zip(batch_entities, batch_embeddings))
+
+                # 批量写入数据库
+                session.execute_write(_batch_set_embeddings, entity_embedding_pairs)
 
             # 数据添加完成后保存图信息
             self.save_graph_info()
 
-    def jsonl_file_add_entity(self, file_path, kgdb_name='neo4j'):
+    async def jsonl_file_add_entity(self, file_path, kgdb_name='neo4j'):
         self.status = "processing"
         kgdb_name = kgdb_name or 'neo4j'
         self.use_database(kgdb_name)  # 切换到指定数据库
@@ -176,7 +227,7 @@ class GraphDatabase:
 
         triples = list(read_triples(file_path))
 
-        self.txt_add_vector_entity(triples, kgdb_name)
+        await self.txt_add_vector_entity(triples, kgdb_name)
 
         self.status = "open"
         # 更新并保存图数据库信息
@@ -355,9 +406,23 @@ class GraphDatabase:
         with self.driver.session() as session:
             return session.execute_read(query, node_name, hops)
 
+    async def aget_embedding(self, text):
+        from src import knowledge_base
+
+        if isinstance(text, list):
+            outputs = await knowledge_base.embed_model.abatch_encode(text, batch_size=40)
+            return outputs
+        else:
+            outputs = await knowledge_base.embed_model.aencode(text)
+            return outputs
+
     def get_embedding(self, text):
-        with torch.no_grad():
-            from src import knowledge_base
+        from src import knowledge_base
+
+        if isinstance(text, list):
+            outputs = knowledge_base.embed_model.batch_encode(text, batch_size=40)
+            return outputs
+        else:
             outputs = knowledge_base.embed_model.encode([text])[0]
             return outputs
 

src/models/embedding.py

@@ -1,6 +1,7 @@
 import os
 import json
 import requests
+import asyncio
 from FlagEmbedding import FlagModel
 from zhipuai import ZhipuAI
 
@@ -26,6 +27,15 @@ class BaseEmbeddingModel:
     def encode_queries(self, queries):
         return self.predict(queries)
 
+    async def aencode(self, message):
+        return await asyncio.to_thread(self.encode, message)
+
+    async def aencode_queries(self, queries):
+        return await asyncio.to_thread(self.encode_queries, queries)
+
+    async def abatch_encode(self, messages, batch_size=20):
+        return await asyncio.to_thread(self.batch_encode, messages, batch_size)
+
     def batch_encode(self, messages, batch_size=20):
         logger.info(f"Batch encoding {len(messages)} messages")
         data = []

web/src/views/GraphView.vue

@@ -198,10 +198,15 @@ const addDocumentByFile = () => {
   })
   .then(response => response.json())
   .then((data) => {
-    message.success(data.message);
-    state.showModal = false;
+    if (data.status === 'success') {
+      message.success(data.message);
+      state.showModal = false;
+    } else {
+      throw new Error(data.message);
+    }
   })
   .catch((error) => {
+    console.error(error)
     message.error(error.message);
   })
   .finally(() => state.precessing = false)