Notebooks
M
Meta Llama Knowledge Graphs With Structured Outputs

Knowledge Graphs With Structured Outputs

llamaAIvllmmachine-learning3p-integrationstogetheraillama2LLMllama-cookbookPythonfinetuningpytorchlangchain
alph-notebooks/llama-cookbook / knowledge_graphs_with_structured_outputs.ipynb

Generating Knowledge Graphs with LLMs and Structured Outputs

Open In Colab

Introduction

In this notebook we show how to generate and visualize knowledge graphs using Large Language Models (LLMs), and how to use JSON mode for structured output generation. We will use Together AI JSON mode to access models that will generate JSON structures that represent a knowledge graph (including nodes, edges, labels) for us.

We will then pass these knowledge graph components to GraphViz that will help us generate and visualize these graphs.

Install relevant libraries

[ ]
[2]

Define Knowledge Graph Schema in Pydantic

We need a way of telling the LLM what structure a knowledge graph has - including what information the nodes, edges must have and how they come together to create the overall graph structure. We will do this using pydantic models.

Below we define the required classes.

  • Each node in our graph needs to have an id and a label. The Node class specifies this.
  • Each edge in our graph needs to have a label and must also connect two nodes from the source to destination direction. The Edge class specified this.
  • Each Knowledge graph is a combination of multiple Node and Edge elements. The KnowledgeGraph class specifies this.
[3]
[5]
{'$defs': {'Edge': {'properties': {'source': {'title': 'Source',
,     'type': 'integer'},
,    'target': {'title': 'Target', 'type': 'integer'},
,    'label': {'title': 'Label', 'type': 'string'}},
,   'required': ['source', 'target', 'label'],
,   'title': 'Edge',
,   'type': 'object'},
,  'Node': {'properties': {'id': {'title': 'Id', 'type': 'integer'},
,    'label': {'title': 'Label', 'type': 'string'}},
,   'required': ['id', 'label'],
,   'title': 'Node',
,   'type': 'object'}},
, 'properties': {'nodes': {'items': {'$ref': '#/$defs/Node'},
,   'title': 'Nodes',
,   'type': 'array'},
,  'edges': {'items': {'$ref': '#/$defs/Edge'},
,   'title': 'Edges',
,   'type': 'array'}},
, 'title': 'KnowledgeGraph',
, 'type': 'object'}

Function to Generate a Knowledge Graph using LLMs

Here we use Together AI to access Meta-Llama-3.1-70B in strict JSON mode. We pass in instructions to get the LLM to generate a knowledge that captures the information about a passed in input string.

We pass in the above KnowledgeGraph JSON schema to guide the structured output generation.

[6]
[7]
[8]
{'nodes': [{'id': 1, 'label': 'Transformer'},
,  {'id': 2, 'label': 'Attention Mechanism'},
,  {'id': 3, 'label': 'Encoder-Decoder Architecture'},
,  {'id': 4, 'label': 'Self-Attention'},
,  {'id': 5, 'label': 'Vaswani et al.'},
,  {'id': 6, 'label': 'Sequence-to-Sequence Models'},
,  {'id': 7, 'label': 'Recurrent Neural Networks (RNNs)'},
,  {'id': 8, 'label': 'Convolutional Neural Networks (CNNs)'},
,  {'id': 9, 'label': 'Parallelization'}],
, 'edges': [{'source': 1, 'target': 2, 'label': 'utilizes'},
,  {'source': 1, 'target': 3, 'label': 'employs'},
,  {'source': 2, 'target': 4, 'label': 'includes'},
,  {'source': 3, 'target': 6, 'label': 'enhances'},
,  {'source': 5, 'target': 1, 'label': 'proposes'},
,  {'source': 6, 'target': 7, 'label': 'replaces'},
,  {'source': 6, 'target': 8, 'label': 'outperforms'},
,  {'source': 1, 'target': 9, 'label': 'enables'}]}

Use GraphViz to Create and Visualize the Graph

Below we write a function that takes the knowledge graph JSON structure above and creates a GraphViz Directed Graph(Digraph) and generates a picture of the graph.

[9]
[10]
Graph rendered and saved to knowledge_graph.png
[11]
Output

We were able to get a LLM to generate a knowledge graph by telling is the structure that a knowledge graph library, GraphViz, expects the output to be in. We then used GraphViz to create and visualize the knowledge graph above!

Learn more about how to use JSON mode in the docs here!