From 61ba90d301cffea32317e3efb18600731e865d5d Mon Sep 17 00:00:00 2001
From: QiyiJiang <3442721524@qq.com>
Date: Wed, 24 Jul 2024 12:23:31 +0000
Subject: [PATCH] =?UTF-8?q?query=5Fby=5Fvector=E6=B7=BB=E5=8A=A0=E6=A3=80?=
 =?UTF-8?q?=E7=B4=A2=E6=95=B0=E9=87=8F=E5=92=8C=E9=98=88=E5=80=BC=E4=B8=A4?=
 =?UTF-8?q?=E4=B8=AA=E5=8F=82=E6=95=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/core/graphbase.py | 20 +++++++++++++++++---
 1 file changed, 17 insertions(+), 3 deletions(-)

diff --git a/src/core/graphbase.py b/src/core/graphbase.py
index 39fe393f..1906e009 100644
--- a/src/core/graphbase.py
+++ b/src/core/graphbase.py
@@ -5,10 +5,15 @@ from neo4j import GraphDatabase as GD
 # from plugins import pdf2txt, OneKE
 from transformers import AutoTokenizer, AutoModel
 from FlagEmbedding import FlagModel, FlagReranker
+import warnings
 
 from plugins import pdf2txt
 from plugins.oneke import OneKE
 
+warnings.filterwarnings("ignore", category=UserWarning)
+
+
+
 UIE_MODEL = None
 
 class GraphDatabase:
@@ -272,10 +277,19 @@ class GraphDatabase:
         with self.driver.session() as session:
             return session.execute_read(query, keyword)
 
-    def query_by_vector(self, entity_name, kgdb_name='neo4j', hops=2):
+    def query_by_vector(self, entity_name,  num_of_res=2, threshold=0.9,kgdb_name='neo4j', hops=2):
         self.use_database(kgdb_name)
         result = self.query_by_vector_tep(entity_name)
-        ans = self.query_specific_entity(result[0][0], hops) # 这里是只获取第一个 TODO: 优化
+        querys = []
+        for i in range(num_of_res):
+            if result[i][1] > threshold:
+                querys.append(result[i][0])
+            else:
+                break
+        ans = []
+        for query in querys:
+            tep = self.query_specific_entity(query, hops) # 这里是只获取第一个 TODO: 优化
+            ans.extend(tep) 
         return ans
 
     def query_node_info(self, node_name, kgdb_name='neo4j', hops = 2):
@@ -354,7 +368,7 @@ if __name__ == "__main__":
     #         }
     #         ]
 
-    # print(kgdb.query_by_vector("维C"))
+    # kgdb.query_by_vector("z")
     # def format_query_results(results):
     #     formatted_results = {"nodes": [], "edges": []}