forked from langchain-ai/langchain
-
Notifications
You must be signed in to change notification settings - Fork 0
/
ingest.py
47 lines (40 loc) · 1.37 KB
/
ingest.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
from pathlib import Path
from langchain.text_splitter import TokenTextSplitter
from langchain_community.document_loaders import TextLoader
from langchain_community.embeddings.openai import OpenAIEmbeddings
from langchain_community.graphs import Neo4jGraph
from langchain_community.vectorstores import Neo4jVector
txt_path = Path(__file__).parent / "dune.txt"
graph = Neo4jGraph()
# Load the text file
loader = TextLoader(str(txt_path))
documents = loader.load()
# Define chunking strategy
parent_splitter = TokenTextSplitter(chunk_size=512, chunk_overlap=24)
child_splitter = TokenTextSplitter(chunk_size=100, chunk_overlap=24)
# Store parent-child patterns into graph
parent_documents = parent_splitter.split_documents(documents)
for parent in parent_documents:
child_documents = child_splitter.split_documents([parent])
params = {
"parent": parent.page_content,
"children": [c.page_content for c in child_documents],
}
graph.query(
"""
CREATE (p:Parent {text: $parent})
WITH p
UNWIND $children AS child
CREATE (c:Child {text: child})
CREATE (c)-[:HAS_PARENT]->(p)
""",
params,
)
# Calculate embedding values on the child nodes
Neo4jVector.from_existing_graph(
OpenAIEmbeddings(),
index_name="retrieval",
node_label="Child",
text_node_properties=["text"],
embedding_node_property="embedding",
)