Spaces:
Sleeping
Sleeping
| import json | |
| import pydot | |
| from sentence_transformers import SentenceTransformer, util | |
| import numpy as np | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| import torch | |
| import openai | |
| from openai import OpenAI | |
| from PIL import Image | |
| system_text = """You are an expert AI that extracts knowledge graphs from text and outputs JSON files with the extracted knowledge, and nothing more. Here's how the JSON is broken down. | |
| Entity dictionaries are organized in a list | |
| Every entity mentioned in the text has its own entity dictionary, in which the name of the entity is the key, and the value is a list of relationships. | |
| Each relationship contains a short word or two accurately describing the relationship to the other entity as the key, and then the other entity as a value. | |
| All inverses of these relationships are represented in the relationship list of the other entities. This is REALY IMPORTANT. For example if Apple created the iPhone, it is also important to note that the iPhone was created by Apple (each entity should have this relsationship from their perspective). | |
| Non specified relationships are also inferred (if person X is the son of person Y, and person Z is person X's sibling, person Z is also the child of person Y). | |
| The JSON contains NO NEW LINES. All the data should be on one line. | |
| Every entity has a "description" relationship which provides a short description of what it is in a few words. If the description references another entity, then this relationship MUST be graphed, even if it is redundant. | |
| Relationships are only created about facts, not just any connection between two entities mentioned in the text. | |
| Example output: | |
| [{"Toki Pona": [{"description": "philosophical artistic constructed language"}, {"translated as": "the language of good"}, {"created by": "Sonja Lang"}, {"first published": "2001"}, {"complete form published in": "Toki Pona: The Language of Good"}, {"supplementary dictionary": "Toki Pona Dictionary"}], "Sonja Lang": [{"description": "Canadian linguist and translator"}, {"creator of": "Toki Pona"}], "Toki Pona: The Language of Good": [{"description": "book"}, {"published in": "2014"}, {"language": "Toki Pona"}], "Toki Pona Dictionary": [{"description": "dictionary"}, {"released in": "July 2021"}, {"based on": "community usage"}]}]""" | |
| class KnowledgeGraph: | |
| def __init__(self,api_key,kg_file=""): | |
| openai.api_key = api_key | |
| self.graph = pydot.Dot(graph_type="digraph") | |
| self.entities = {} | |
| self.fact_scores = {} | |
| self.device = "cuda" if torch.cuda.is_available() else "cpu" | |
| self.model = SentenceTransformer('all-MiniLM-L6-v2').to(self.device) | |
| self.entity_embeddings = {} | |
| if kg_file!="": | |
| self.load_graph(kg_file) | |
| def add_entity(self, name, description): | |
| if name not in self.entities: | |
| self.entities[name] = {"description": description} | |
| entity_node = pydot.Node(name, label=f"{name}\n({description})") | |
| self.graph.add_node(entity_node) | |
| self.entity_embeddings[name] = self.model.encode(name)#+": \n"+"\n".join([key+": "+kg.entities[name][key] for key in kg.entities[name]])) | |
| print("added embedding") | |
| def add_relationship(self, entity1, relationship, entity2): | |
| if entity1 in self.entities: | |
| try: | |
| self.entities[entity1][relationship] += ", "+entity2 | |
| except: | |
| self.entities[entity1][relationship] = entity2 | |
| edge = pydot.Edge(entity1, entity2, label=relationship) | |
| self.graph.add_edge(edge) | |
| def update_graph(self, json_str,clean=True): | |
| try: | |
| data = json.loads(json_str) | |
| except: | |
| print("GPT4 failed to create a valid JSON. Input may be too long for processing.") | |
| return | |
| for entity_dict in data: | |
| for entity, relationships in entity_dict.items(): | |
| try: | |
| self.add_entity(entity, relationships[0]["description"]) | |
| except: | |
| self.add_entity(entity, "") | |
| for rel in relationships[1:]: | |
| for relationship, other_entity in rel.items(): | |
| try: | |
| self.add_relationship(entity, relationship, other_entity) | |
| except: | |
| for o in other_entity: | |
| self.add_relationship(entity, relationship, o) | |
| if clean: | |
| for entity_dict in data: | |
| for entity, relationships in entity_dict.items(): | |
| self.clean_graph(entity) | |
| def display_graph(self, output_file="knowledge_graph.png"): | |
| self.graph.write_png(output_file) | |
| img = Image.open(output_file) | |
| img.show() | |
| return img | |
| def search(self, query, n=5): | |
| if len(self.entity_embeddings)<5: | |
| n = len(self.entity_embeddings) | |
| query_embedding = self.model.encode(query) | |
| query_tensor = torch.tensor([query_embedding]) | |
| entity_tensor = torch.tensor(list(self.entity_embeddings.values())) | |
| similarities = util.cos_sim(query_tensor, entity_tensor).numpy() | |
| top_indices = np.argsort(similarities[0])[-n:][::-1] | |
| results = [(list(self.entity_embeddings.keys())[index], similarities[0][index]) for index in top_indices] | |
| return results | |
| def related_entities(self,query, n=5): | |
| query_embedding = self.model.encode(query) | |
| query_tensor = torch.tensor([query_embedding]) | |
| potentities = [key+": "+self.entities[key]["description"] for key in self.entities] | |
| entity_tensor = self.model.encode(potentities) | |
| similarities = util.cos_sim(query_tensor, entity_tensor).numpy() | |
| if len(similarities)<n: | |
| n = len(similarities) | |
| top_indices = np.argsort(similarities[0])[-n:][::-1] | |
| results = [potentities[index] for index in top_indices] | |
| return results | |
| def text_to_data(self,text): | |
| system = {"role":"system","content":system_text} | |
| messages = [system] | |
| try: | |
| related = self.related_entities(text) | |
| text = text+f"\n\nGenerate the JSON for the text above, remembering to add inverse relationships and inferences. Here are some related entities already in the graph. If you are adding information about any of them, refer to them by the names below (otherwise ignore this information):\n\n{str(related)}" | |
| except: | |
| pass | |
| messages.append({"role":"user","content":text}) | |
| output = openai.ChatCompletion.create(model="gpt-4o",messages=messages)["choices"][0]["message"].to_dict()["content"] | |
| return output | |
| def learn(self,text,show_output=False): | |
| json_str = self.text_to_data(text) | |
| if show_output: | |
| print(json_str) | |
| self.update_graph(json_str) | |
| def graph_search(self,query,n=5,path="subgraph.png"): | |
| results = self.search(query, n) | |
| if len(results)<n: | |
| n = len(results) | |
| top_ents = [results[i][0] for i in range(n)] | |
| data = [{ent:[{key:self.entities[ent][key]} for key in self.entities[ent]]} for ent in top_ents] | |
| new = KnowledgeGraph() | |
| json_string = json.dumps(data) | |
| new.update_graph(str(json_string),clean=False) | |
| new.display_graph(path) | |
| def text_search(self,query,n=3): | |
| results = self.search(query, n) | |
| keys = [r[0] for r in results] | |
| potentities = [key+": "+str(self.entities[key]) for key in keys] | |
| for p in potentities: | |
| print(p) | |
| def qa_search(self,query,n=5): | |
| results = self.search(query, n) | |
| 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]] | |
| 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] | |
| 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.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}" | |
| messages.append({"role":"user","content":text}) | |
| output = openai.ChatCompletion.create(model="gpt-3.5-turbo",messages=messages)["choices"][0]["message"].to_dict()["content"] | |
| return output | |
| def clean_graph(self,key): | |
| facts = [key+": "+str(rel).replace("description","is")+" "+str(self.entities[key][rel]) for rel in self.entities[key]] | |
| rels = [rel for rel in self.entities[key]] | |
| fact_embs = self.model.encode(facts) | |
| scores = util.cos_sim(fact_embs,fact_embs) | |
| pairs = [] | |
| for i in range(len(scores)): | |
| for j in range(len(scores[i])): | |
| if round(scores[i][j].item(),3)!=1.0 and scores[i][j]>0.7: | |
| if (facts[i],facts[j]) not in pairs and (facts[j],facts[i]) not in pairs: | |
| pairs.append((facts[i],facts[j])) | |
| for pair in pairs: | |
| system = {"role":"system","content":"You are a helpful chatbot that only outputs YES or NO"} | |
| messages = [system] | |
| messages.append({"role":"user","content":f"Do these two facts in our database express the same thing?: {pair}"}) | |
| output = openai.ChatCompletion.create(model="gpt-4o",messages=messages)["choices"][0]["message"].to_dict()["content"] | |
| if "yes" in output.lower(): | |
| bad_index = facts.index(pair[1]) | |
| redundant = rels[bad_index] | |
| del self.entities[key][redundant] | |
| good_index = facts.index(pair[0]) | |
| validated = rels[bad_index] | |
| try: | |
| self.fact_scores[(key,validated)]+=1 | |
| except: | |
| self.fact_scores[(key,validated)]=1 | |
| def load_graph(self,kg_file): | |
| with open(kg_file) as f: | |
| lines = f.readlines() | |
| graph_data = "\n".join(lines[:-1]) | |
| ents = eval(lines[-1]) | |
| data = [{ent:[{key:ents[ent][key]} for key in ents[ent]]} for ent in ents] | |
| json_string = json.dumps(data) | |
| print(json_string) | |
| self.update_graph(str(json_string)) | |
| self.graph = pydot.graph_from_dot_data(graph_data)[0] | |
| def save_graph(self,filename="mygraph.kg"): | |
| with open(filename,"w") as f: | |
| f.write("") | |
| self.graph.write_dot(filename) | |
| with open(filename,"a") as f: | |
| f.write("\n") | |
| f.write(str(self.entities)) |