query_by_vector添加检索数量和阈值两个参数

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": []}