txtai is an all-in-one embeddings database for semantic search, LLM orchestration and language model workflows.
Embeddings databases are a union of vector indexes (sparse and dense), graph networks and relational databases. This enables vector search with SQL, topic modeling, retrieval augmented generation and more.
Up until txtai 7.0, semantic graphs only supported automatic relationship detection. Now relationships can be loaded directly into a txtai database. This article will demonstrate how to work with this new feature.
Install dependencies
Install txtai and all dependencies.
# Install txtai
pip install txtai[graph] autoawq
Load Wikipedia database
The txtai-wikipedia database stores all Wikipedia article abstracts as of January 2024. This database is a great way to explore a wide variety of topics. It also has the number of page views integrated in, which enables pulling frequently viewed or popular articles on a topic.
For this example, we'll work with a couple articles related to Viking raids in France.
from txtai import Embeddings
# Load dataset
wikipedia = Embeddings()
wikipedia.load(provider="huggingface-hub", container="neuml/txtai-wikipedia")
query = """
SELECT id, text FROM txtai WHERE similar('Viking raids in France') and percentile >= 0.5
"""
results = wikipedia.search(query, 5)
for x in results:
print(x)
LLM-driven entity extraction
Now that we have a couple relevant articles, let's go through and run an entity extraction process. For this task, we'll have a LLM prompt to do the work.
import json
from txtai import LLM
# Load LLM
llm = LLM("TheBloke/Mistral-7B-OpenOrca-AWQ")
data = []
for result in results:
prompt = f"""<|im_start|>system
You are a friendly assistant. You answer questions from users.<|im_end|>
<|im_start|>user
Extract an entity relationship graph from the following text. Output as JSON
Nodes must have label and type attributes. Edges must have source, target and relationship attributes.
text: {result['text']} <|im_end|>
<|im_start|>assistant
"""
try:
data.append(json.loads(llm(prompt, maxlength=4096)))
except:
pass
Build an embeddings database
Now that we've extracted entities from the documents, next we'll review and load a graph network using these entity-relationships.
def stream():
nodes = {}
for row in data.copy():
# Create nodes
for node in row["nodes"]:
if node["label"] not in nodes:
node["id"] = len(nodes)
nodes[node["label"]] = node
for edge in row["edges"]:
source = nodes.get(edge["source"])
target = nodes.get(edge["target"])
if source and target:
if "relationships" not in source:
source["relationships"] = []
source["relationships"].append({"id": target["id"], "relationship": edge["relationship"]})
return nodes.values()
# Create embeddings instance with a semantic graph
embeddings = Embeddings(
autoid = "uuid5",
path = "intfloat/e5-base",
instructions = {
"query": "query: ",
"data": "passage: "
},
columns = {
"text": "label"
},
content = True,
graph = {
"approximate": False,
"topics": {}
}
)
embeddings.index(stream())
Show the network
Now let's visualize the entity-relationship network. This period might not be that familar to most, unless you've watched the Vikings TV series, in which case it should make sense.
import matplotlib.pyplot as plt
import networkx as nx
def plot(graph):
labels = {x: f"{graph.attribute(x, 'text')} ({x})" for x in graph.scan()}
lookup = {
"Person": "#d32f2f",
"Location": "#0277bd",
"Event": "#e64980",
"Role": "#757575"
}
colors = []
for x in graph.scan():
value = embeddings.search("select type from txtai where id = :x", parameters={"x": x})[0]["type"]
colors.append(lookup.get(value, "#7e57c2"))
options = {
"node_size": 2000,
"node_color": colors,
"edge_color": "#454545",
"font_color": "#efefef",
"font_size": 11,
"alpha": 1.0
}
fig, ax = plt.subplots(figsize=(20, 9))
pos = nx.spring_layout(graph.backend, seed=0, k=3, iterations=250)
nx.draw_networkx(graph.backend, pos=pos, labels=labels, **options)
ax.set_facecolor("#303030")
ax.axis("off")
fig.set_facecolor("#303030")
plot(embeddings.graph)
Graph traversal
The last thing we'll cover is extracting a specific path from the graph. Let's show a path from node 8 to node 5 requiring a specific relationship type to start.
A new feature of txtai 7.0 is the ability to return a graph of search results. This is a powerful addition as not only do we get the search results but we get how the search results relate to each other.
# Traverse graph looking for certain nodes and edge values
g = embeddings.graph.search("""
MATCH P=(A{id: 8})-[R1]->()-[*1..3]->(D{id:5})
WHERE
R1.relationship == "has_object"
RETURN P
""", graph=True)
plot(g)
Wrapping up
This article showed how knowledge graphs can be created with LLM-driven entity extraction. Automatically derived relationships via semantic similarity and manually specified ones are a powerful combination!
Top comments (0)