|
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() |
|
|
|
|
|
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])}") |
|
|
|
|
|
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) |
|
|
|
|
|
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() |
