Update GPT4KG.py

GPT4KG.py

@@ -150,9 +150,37 @@ class KnowledgeGraph:
       top_indices = np.argsort(similarities[0])[-n:][::-1]
       results = [facts[index] for index in top_indices]
       return results
+
+    def fact_search(self,query,facts,n=5):
+      query_embedding = self.model.encode(query)
+      query_tensor = torch.tensor([query_embedding])
+      fact_tensor = self.model.encode(facts)
+      similarities = util.cos_sim(query_tensor, fact_tensor).numpy()
+      if len(similarities[0])<n:
+        n = len(similarities)
+      top_indices = np.argsort(similarities[0])[-n:][::-1]
+      return top_indices
+
+    def branch_search(self,query,num_branches=2,window_size=5):
+      results = self.search(query, window_size)
+      keys = [r[0] for r in results]
+      facts = [key+": "+str(rel).replace("description","is")+" "+str(self.entities[key][rel]) for key in keys for rel in self.entities[key]]
+      top_indices = self.fact_search(query,facts)
+      next_ents = [str(self.entities[key][rel]) for key in keys for rel in self.entities[key]]
+      top_ents = [next_ents[index] for index in top_indices]
+      if num_branches == 0:
+        return [facts[index] for index in top_indices]
+      for b in range(num_branches):
+        next_facts = [facts[i]+". "+keys[i]+": "+str(rel).replace("description","is")+" "+str(self.entities[keys[i]][rel]) for i in range(len(keys)) for rel in self.entities[keys[i]] if facts[i].find(keys[i])!=0]
+        facts = next_facts
+        top_indices = self.fact_search(query,facts)
+        answers = [facts[index] for index in top_indices]
+        next_ents = [str(self.entities[key][rel]) for key in keys for rel in self.entities[key]]
+        top_ents = [next_ents[index] for index in top_indices]
+      return answers
     
     def chat_qa(self,query):
-      results = self.qa_search(query)
+      results = self.branch_search(query)
       system = {"role":"system","content":"You are a helpful chatbot that answers questions based on data in your fact database."}
       messages = [system]
       text = f"Question: {query}\n\nFact Data: \n{results}"