Create insertData.py

insertData.py

@@ -0,0 +1,130 @@
+from langchain_community.graphs.neo4j_graph import Neo4jGraph
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.document_loaders import WikipediaLoader
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+
+url = "neo4j+s://ddb8863b.databases.neo4j.io"
+username = "neo4j"
+password = "vz6OLij_IrY-cSIgSMhUWxblTUzH8m4bZaBeJGgmtU0"
+graph = Neo4jGraph(url=url, username=username, password=password)
+
+
+text_splitter = RecursiveCharacterTextSplitter(
+    chunk_size=512,
+    length_function=len, 
+    is_separator_regex=False,)
+
+query = "Dune (Frank Herbert)"
+raw_documents = WikipediaLoader(query=query).load_and_split(text_splitter=text_splitter)
+
+
+tokenizer = AutoTokenizer.from_pretrained("Babelscape/rebel-large")
+model = AutoModelForSeq2SeqLM.from_pretrained("Babelscape/rebel-large")
+
+def extract_relations_from_model_output(text):
+    relations = []    
+    relation, subject, relation, object_ = '', '', '', ''    
+    text = text.strip()    
+    current = 'x'    
+    text_replaced = text.replace("", "").replace("", "").replace("", "")    
+    for token in text_replaced.split():
+        if token == "":
+            current = 't' 
+        if relation != '': 
+            relations.append({ 
+                'head': subject.strip(),
+                'type': relation.strip(),
+                        'tail': object_.strip()
+                })   
+            relation = ''
+            subject = ''
+        elif token == "":
+                current = 's' 
+                if relation != '':
+                    relations.append({ 
+                        'head': subject.strip(), 
+                        'type': relation.strip(),  
+                        'tail': object_.strip()
+                        })   
+                    object_ = '' 
+        elif token == "": 
+                current = 'o'
+                relation = ''
+        else:   
+                if current == 't':
+                    subject += ' ' + token  
+                elif current == 's':
+                    object_ += ' ' + token 
+                elif current == 'o': 
+                    relation += ' ' + token
+    if subject != '' and relation != '' and object_ != '':    
+            relations.append({
+                'head': subject.strip(), 
+                'type': relation.strip(), 
+                'tail': object_.strip()  
+                    }) 
+    return relations
+
+class KB():
+    def __init__(self):
+        self.relations = []
+    def are_relations_equal(self, r1, r2):
+        return all(r1[attr] == r2[attr] for attr in ["head", "type", "tail"])
+    def exists_relation(self, r1): 
+        return any(self.are_relations_equal(r1, r2) for r2 in self.relations)
+    def add_relation(self, r):  
+        if not self.exists_relation(r): 
+            self.relations.append(r) 
+    def print(self):  
+        print("Relations:") 
+        for r in self.relations: 
+            print(f"  {r}")
+
+def from_small_text_to_kb(text, verbose=False):
+    kb = KB()  
+
+    # Tokenizer text
+    model_inputs = tokenizer(text, max_length=512, padding=True, truncation=True,return_tensors='pt')
+
+    if verbose: 
+        print(f"Num tokens: {len(model_inputs['input_ids'][0])}")
+        
+        # Generate   
+    gen_kwargs = { 
+        "max_length": 216,
+        "length_penalty": 0, 
+        "num_beams": 3,  
+        "num_return_sequences": 3
+        } 
+        
+    generated_tokens = model.generate( 
+        **model_inputs, 
+        **gen_kwargs, 
+        )  
+        
+    decoded_preds = tokenizer.batch_decode(generated_tokens, skip_special_tokens=False)  
+
+    # create kb 
+    for sentence_pred in decoded_preds: 
+        relations = extract_relations_from_model_output(sentence_pred) 
+        for r in relations: 
+            kb.add_relation(r) 
+        
+    return kb
+
+for doc in raw_documents:
+    kb = from_small_text_to_kb(doc.page_content, verbose=True)
+
+
+    for relation in kb.relations:
+        head = relation['head']
+        relationship = relation['type']
+        tail = relation['tail']
+
+
+        cypher = f"MERGE (h:`{head}`)" + f" MERGE (t:`{tail}`)" + f" MERGE (h)-[:`{relationship}`]->(t)"
+        print(cypher)
+        graph.query(cypher)
+
+
+graph.refresh_schema()