xbe.at

LLMs and generative AI / RAG and retrieval

Enhancing RAG with Knowledge Graph in Python

7 min read, 14 sections

Contents
  1. Introduction to Knowledge Graphs
  2. What is Retrieval-Augmented Generation (RAG)?
  3. Integrating Knowledge Graphs with RAG
  4. Knowledge Graph Construction
  5. Knowledge Graph Embedding
  6. Knowledge Graph Querying with SPARQL
  7. Knowledge Graph Visualization
  8. Knowledge Graph Reasoning
  9. Knowledge Graph Fusion
  10. Knowledge Graph Alignment
  11. Knowledge Graph Completion
  12. Knowledge Graph Update
  13. Knowledge Graph Evaluation
  14. Additional Resources

Introduction to Knowledge Graphs

Knowledge Graphs are structured representations of real-world entities and their relationships. They provide a powerful way to store, organize, and query complex information, making them ideal for enhancing Retrieval-Augmented Generation (RAG) models.

Code:

import networkx as nx

# Create a knowledge graph
kg = nx.DiGraph()

# Add nodes (entities)
kg.add_node("Python", type="Programming Language")
kg.add_node("Java", type="Programming Language")
kg.add_node("C++", type="Programming Language")

# Add edges (relationships)
kg.add_edge("Python", "Java", relation="similar_to")
kg.add_edge("Python", "C++", relation="similar_to")

What is Retrieval-Augmented Generation (RAG)?

Retrieval-Augmented Generation (RAG) is a natural language processing approach that combines retrieval and generation models. It first retrieves relevant information from a knowledge source (e.g., a knowledge graph) and then generates a response based on the retrieved information.

Code:

from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration

tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
retriever = RagRetriever.from_pretrained("facebook/rag-token-nq", index_name="nq-open", passages=True)
model = RagSequenceForGeneration.from_pretrained("facebook/rag-token-nq")

question = "What is the capital of France?"
inputs = tokenizer(question, return_tensors="pt")
outputs = model.generate(**inputs, max_length=200, num_beams=2, early_stopping=True)
answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(answer)

Integrating Knowledge Graphs with RAG

To enhance RAG models with knowledge graphs, we need to create a knowledge retriever that can retrieve relevant information from the knowledge graph based on the input query. This retriever can then be incorporated into the RAG pipeline.

Code:

from rdflib import Graph

# Load the knowledge graph
kg = Graph().parse("path/to/knowledge_graph.ttl", format="turtle")

def retrieve_from_kg(query):
    # Define SPARQL query based on the input query
    sparql_query = """
        PREFIX : <http://example.org/>
        SELECT ?subject ?predicate ?object
        WHERE {
            ?subject ?predicate ?object .
            FILTER (
                regex(?subject, "%s", "i")
                || regex(?predicate, "%s", "i")
                || regex(?object, "%s", "i")
            )
        }
    """ % (query, query, query)

    # Execute the SPARQL query
    results = kg.query(sparql_query)

    # Return the retrieved triples
    return [(str(row.subject), str(row.predicate), str(row.object)) for row in results]

Knowledge Graph Construction

Before integrating a knowledge graph with RAG, you need to construct the knowledge graph. This process involves extracting entities, relations, and other relevant information from various data sources and structuring them in a graph format.

Code:

from rdflib import Graph, Literal, Namespace, URIRef

# Define namespaces
kg_ns = Namespace("http://example.org/kg#")

# Create a new knowledge graph
kg = Graph()

# Add triples to the knowledge graph
kg.add((URIRef(kg_ns["Paris"]), URIRef(kg_ns["capitalOf"]), URIRef(kg_ns["France"])))
kg.add((URIRef(kg_ns["Python"]), URIRef(kg_ns["programmingLanguage"]), Literal("Python")))
kg.add((URIRef(kg_ns["Java"]), URIRef(kg_ns["programmingLanguage"]), Literal("Java")))

# Serialize the knowledge graph to a file
kg.serialize("path/to/knowledge_graph.ttl", format="turtle")

Knowledge Graph Embedding

Knowledge graph embedding is the process of representing entities and relations in a knowledge graph as dense vector representations. This can improve the performance of knowledge retrieval and enhance the integration of knowledge graphs with RAG models.

Code:

import numpy as np
from ampligraph.datasets import load_from_rdf
from ampligraph.latent_features import TransE

# Load the knowledge graph
kg = load_from_rdf("path/to/knowledge_graph.ttl", "turtle")

# Define the TransE model
model = TransE(batches_count=64, seed=0, epochs=200, k=100, eta=20)

# Train the model
X = np.array([kg.train_idx_obt[:]])
model.fit(X)

# Get embeddings for entities and relations
entity_embeddings = model.ent_embeddings
relation_embeddings = model.rel_embeddings

Knowledge Graph Querying with SPARQL

SPARQL (SPARQL Protocol and RDF Query Language) is a standard query language for querying and manipulating data stored in RDF format, which is commonly used for representing knowledge graphs. SPARQL allows you to retrieve specific information from the knowledge graph based on your queries.

Code:

from rdflib import Graph

# Load the knowledge graph
kg = Graph().parse("path/to/knowledge_graph.ttl", format="turtle")

# Define a SPARQL query
query = """
    PREFIX : <http://example.org/>
    SELECT ?capital ?country
    WHERE {
        ?capital :capitalOf ?country .
    }
"""

# Execute the SPARQL query
results = kg.query(query)

# Print the results
for row in results:
    print(f"{row.capital.value} is the capital of {row.country.value}")

Knowledge Graph Visualization

Visualizing knowledge graphs can provide insights into the structure, relationships, and patterns within the data. This can be useful for understanding, exploring, and debugging knowledge graphs.

Code:

import networkx as nx
import matplotlib.pyplot as plt

# Load the knowledge graph
kg = nx.read_gml("path/to/knowledge_graph.gml")

# Draw the knowledge graph
pos = nx.spring_layout(kg)
nx.draw(kg, pos, with_labels=True, node_color="skyblue", edge_color="gray")
plt.axis("off")
plt.show()

Knowledge Graph Reasoning

Knowledge graph reasoning involves inferring new knowledge from existing knowledge in the graph. This can be achieved through various techniques, such as rule-based reasoning, statistical relational learning, or deep learning-based methods.

Code:

from ampligraph.latent_features import ComplEx
from ampligraph.utils import create_bf

# Load the knowledge graph
kg = load_from_rdf("path/to/knowledge_graph.ttl", "turtle")

# Define the ComplEx model
model = ComplEx(batches_count=64, seed=0, epochs=200, k=100, eta=20)

# Train the model
X = np.array([kg.train_idx_obt[:]])
model.fit(X)

# Create a new batch of triples for inference
new_triples = create_bf(model, 100)

# Infer new knowledge from the knowledge graph
model.predict(new_triples)

Knowledge Graph Fusion

Knowledge graph fusion involves combining multiple knowledge graphs into a single, unified graph. This can be useful when working with diverse data sources or integrating complementary knowledge bases.

Code:

from rdflib import Graph

# Load the first knowledge graph
kg1 = Graph().parse("path/to/kg1.ttl", format="turtle")

# Load the second knowledge graph
kg2 = Graph().parse("path/to/kg2.ttl", format="turtle")

# Create a new graph to hold the fused knowledge graph
fused_kg = Graph()

# Add triples from kg1 to the fused graph
for s, p, o in kg1.triples((None, None, None)):
    fused_kg.add((s, p, o))

# Add triples from kg2 to the fused graph
for s, p, o in kg2.triples((None, None, None)):
    fused_kg.add((s, p, o))

# Optionally, perform deduplication or conflict resolution
# ...

# Serialize the fused knowledge graph to a file
fused_kg.serialize("path/to/fused_kg.ttl", format="turtle")

Knowledge Graph Alignment

Knowledge graph alignment is the process of finding correspondences between entities or relations in different knowledge graphs. This is particularly useful when integrating or fusing multiple knowledge graphs, as it helps identify and resolve potential conflicts or redundancies.

Code:

from ampligraph.evaluation import hits_at_n_train
from ampligraph.latent_features import TransE

# Load the first knowledge graph
kg1 = load_from_rdf("path/to/kg1.ttl", "turtle")

# Load the second knowledge graph
kg2 = load_from_rdf("path/to/kg2.ttl", "turtle")

# Define the TransE model
model = TransE(batches_count=64, seed=0, epochs=200, k=100, eta=20)

# Train the model on kg1
X1 = np.array([kg1.train_idx_obt[:]])
model.fit(X1)

# Evaluate the model on kg2 to find entity alignments
hits = hits_at_n_train(model, kg2, np.array([kg2.train_idx_obt[:]]), hits=[1, 3, 10])
print(hits)

Knowledge Graph Completion

Knowledge graph completion is the task of inferring missing facts or relations in a knowledge graph. This can be achieved through various techniques, such as rule mining, tensor factorization, or neural network-based methods.

Code:

import numpy as np
from ampligraph.latent_features import ComplEx
from ampligraph.utils import create_bf

# Load the knowledge graph
kg = load_from_rdf("path/to/knowledge_graph.ttl", "turtle")

# Define the ComplEx model
model = ComplEx(batches_count=64, seed=0, epochs=200, k=100, eta=20)

# Train the model
X = np.array([kg.train_idx_obt[:]])
model.fit(X)

# Create a batch of triples with missing components
new_triples = create_bf(model, 100, form="SP?")

# Predict the missing components (objects)
predictions = model.predict(new_triples)

Knowledge Graph Update

Knowledge graphs are dynamic and may require updates as new information becomes available. Updating a knowledge graph involves adding, modifying, or removing entities, relations, or facts within the graph.

Code:

from rdflib import Graph, Literal, Namespace, URIRef

# Load the existing knowledge graph
kg = Graph().parse("path/to/knowledge_graph.ttl", format="turtle")

# Define namespaces
kg_ns = Namespace("http://example.org/kg#")

# Add a new triple to the knowledge graph
kg.add((URIRef(kg_ns["Python"]), URIRef(kg_ns["version"]), Literal("3.9")))

# Remove an existing triple from the knowledge graph
kg.remove((URIRef(kg_ns["Java"]), URIRef(kg_ns["programmingLanguage"]), Literal("Java")))

# Modify an existing triple in the knowledge graph
kg.remove((URIRef(kg_ns["Paris"]), URIRef(kg_ns["capitalOf"]), URIRef(kg_ns["France"])))
kg.add((URIRef(kg_ns["Paris"]), URIRef(kg_ns["capitalOf"]), URIRef(kg_ns["France"]), Literal("2024")))

# Serialize the updated knowledge graph to a file
kg.serialize("path/to/updated_kg.ttl", format="turtle")

Knowledge Graph Evaluation

Evaluating the quality and performance of a knowledge graph is essential for ensuring its reliability and effectiveness. Several metrics and techniques can be used for this purpose, such as link prediction, triple classification, and entity resolution.

Code:

from ampligraph.evaluation import hits_at_n_train
from ampligraph.latent_features import TransE

# Load the knowledge graph
kg = load_from_rdf("path/to/knowledge_graph.ttl", "turtle")

# Define the TransE model
model = TransE(batches_count=64, seed=0, epochs=200, k=100, eta=20)

# Train the model
X = np.array([kg.train_idx_obt[:]])
model.fit(X)

# Evaluate the model using link prediction
hits = hits_at_n_train(model, kg, np.array([kg.train_idx_obt[:]]), hits=[1, 3, 10])
print(hits)

Additional Resources

For further learning and exploration of knowledge graphs and their integration with RAG models, here are some additional resources:

These resources provide in-depth information, research papers, and surveys on various aspects of knowledge graphs and their applications in natural language processing tasks, including Retrieval-Augmented Generation (RAG).

More in RAG and retrieval

  1. Fine-tuning T5-small for Retrieval-Augmented Generation (RAG) Using Python3 min

    Introduction to Fine-tuning T5-small for RAG (Retrieval-Augmented Generation)

  2. Enhancing LLM Context with Recursive Summarization Using Python6 min

    Large Language Models (LLMs) have limited context windows. Recursive summarization is a technique to extend this context by iteratively condensing...

  3. GNN-RAG Graph Neural Retrieval for Large Language Model Reasoning in Python5 min

    GNN-RAG (Graph Neural Network for Reasoning and Access to Graphs) is a framework that combines the power of large language models with graph neural networks...

  4. End-to-End Retrieval-Augmented Generation (RAG) Application Using Haystack, Pinecone, and Mistral AI7 min

    End-to-End Retrieval-Augmented Generation (RAG) is an advanced natural language processing technique that combines the power of large language models with...

  5. GraphRAG Graph-based NLP with Python9 min

    GraphRAG is an innovative approach that combines graph-based knowledge representation with Retrieval-Augmented Generation (RAG) for enhanced natural language...

  6. EfficientRAG Supercharging Multi-Hop QA for LLMs12 min

    EfficientRAG is an advanced technique for enhancing multi-hop question answering (QA) in Large Language Models (LLMs) using Python. This method combines...