Constructing a Knowledge Graph from unstructured data with LLMs

Getting unstructured data and chunking it with LangChain

    # Load the OpenAI Wikipedia page
    from langchain.document_loaders import WikipediaLoader
    from langchain.text_splitter import RecursiveCharacterTextSplitter
    raw_documents = WikipediaLoader(query="Chiang Kai-Shek").load()
    
    # Remove unnecessary Wikipedia pages, if need be
    import numpy as np

    for doc in raw_documents:
      print(doc.metadata['source'])

    raw_documents = np.array(raw_documents)
    raw_documents = raw_documents[[0,1,2]]

    # Define chunking strategy
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=2048, chunk_overlap=24)
    
    # Chunk the document
    documents = text_splitter.split_documents(raw_documents)
    for d in documents:
        del d.metadata["summary"]   

The desired JSON structure is predefined and provided to the LLM through an OpenAI function call with LangChain, to restrain the LLM response as a standardized representation of entities and relationships we want extracted from the chunk of text and ready to be ingested into a graph database such as Neo4j.

  from langchain.graphs.graph_document import (
      Node as BaseNode,
      Relationship as BaseRelationship,
      GraphDocument,
  )
  from langchain.schema import Document
  from typing import List, Dict, Any, Optional
  from langchain.pydantic_v1 import Field, BaseModel
  
  class Property(BaseModel):
    """A single property consisting of key and value"""
    key: str = Field(..., description="key")
    value: str = Field(..., description="value")
  
  class Node(BaseNode):
      properties: Optional[List[Property]] = Field(
          None, description="List of node properties")
  
  class Relationship(BaseRelationship):
      properties: Optional[List[Property]] = Field(
          None, description="List of relationship properties"
      )
  
  class KnowledgeGraph(BaseModel):
      """Generate a knowledge graph with entities and relationships."""
      nodes: List[Node] = Field(
          ..., description="List of nodes in the knowledge graph")
      rels: List[Relationship] = Field(
          ..., description="List of relationships in the knowledge graph"
      )

Extracting structured information with OpenAI's gpt-3.5 model

  from langchain.chains.openai_functions import (
      create_structured_output_chain
  )
  from langchain.chat_models import ChatOpenAI
  from langchain.prompts import ChatPromptTemplate
  
  # Initialize LLM
  llm = ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo")
  
  def get_extraction_chain(
      allowed_nodes: Optional[List[str]] = None,
      allowed_rels: Optional[List[str]] = None
      ):
      
      prompt = ChatPromptTemplate.from_messages(
          [(
            "system",
            f"""
  You are a historian extracting information in structured formats to build a knowledge graph around a key historical figure.
  {'- Allowed Node Labels:' + ", ".join(allowed_nodes) if allowed_nodes else ""}
  {'- Allowed Relationship Types:' + ", ".join(allowed_rels) if allowed_rels else ""}
  - Numerical data, like age or other related information, should be incorporated as attributes or properties of the respective nodes.
  - No Separate Nodes for Dates/Numbers: Do not create separate nodes for dates or numerical values. Always attach them as attributes or properties of nodes.
  - Property Format: Properties must be in a key-value format.
  - Quotation Marks: Never use escaped single or double quotes within property values.
  - Naming Convention: Use camelCase for property keys, e.g., "birthDate".
  - Relationship Extraction: Some relationship labels have the same meaning, but may be of different verb tense (e.g., "leader of", "led"), always use the most complete label for that relationship throughout the knowledge graph. 
  - Coreference Resolution: When extracting entities, it's vital to ensure consistency. If an entity, such as "John Doe", is mentioned multiple times in the text but is referred to by different names or pronouns (e.g., "Joe", "he"),
  always use the most complete identifier for that entity throughout the knowledge graph. In this example, use "John Doe" as the entity ID.
            """),
              ("human", "Use the given format to extract information from the following input: {input}"),
              ("human", "Tip: Make sure to answer in the correct format and JSON compliant"),
          ])
      return create_structured_output_chain(KnowledgeGraph, llm, prompt, verbose=True)

Store extracted information into Neo4j

  from langchain.graphs import Neo4jGraph

  url = "bolt://localhost:7687"
  username ="neo4j"
  password = "yourpassword"
  graph = Neo4jGraph(
      url=url,
      username=username,
      password=password,
      database="youdbname"
  )

  def extract_and_store_graph(
    document: Document,
    nodes:Optional[List[str]] = None,
    rels:Optional[List[str]]=None) -> None:
    # Extract graph data using OpenAI functions
    extract_chain = get_extraction_chain(nodes, rels)
    data = extract_chain.run(document.page_content)

    # Construct a graph document
    graph_document = GraphDocument(
      nodes = [map_to_base_node(node) for node in data.nodes],
      relationships = [map_to_base_relationship(rel) for rel in data.rels],
      source = document
    )

    # Store information into a graph
    graph.add_graph_documents([graph_document])

    from tqdm import tqdm

    # Specify which node labels should be extracted by the LLM
    allowed_nodes = ["Person", "Organization", "Location", "Event", "Service", "Object", "Concept"]
    
    for i, d in tqdm(enumerate(documents), total=len(documents)):
      extract_and_store_graph(d, allowed_nodes)

Graph visualization

Additional pre-extraction steps

Summarization

    from langchain.chains import MapReduceDocumentsChain, ReduceDocumentsChain, StuffDocumentsChain

    # Initialize LLM
    llm = ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo")
    
    # Define the map prompt template
    map_prompt = PromptTemplate.from_template(
      "The following is a set of documents: {documents}, please perform summaries of them."
    )
    
    map_chain = LLMChain(llm=llm, prompt=map_prompt)
    
    # Define how to combine these summaries
    reduce_prompt = PromptTemplate.from_template(
      "The following is set of summaries: {documents}, take these and distill it into a final, consolidated summary."
    )

    reduce_chain = LLMChain(llm=llm, prompt=reduce_prompt)

    # Takes the summaries, combines them into a single string, and passes this to an LLMChain
    combine_documents_chain = StuffDocumentsChain(
        llm_chain=reduce_chain,
        document_variable_name="documents"  # This should match the variable name in reduce_prompt
    )
    
    # Combines and iteravely reduces the mapped documents
    reduce_documents_chain = ReduceDocumentsChain(
        combine_documents_chain=combine_documents_chain, # This is final chain that is called.
        token_max=4000 # The maximum number of tokens to group documents into.
    )
    
    # Combining documents by mapping a chain over them, then combining results
    map_reduce_chain = MapReduceDocumentsChain(
        llm_chain=map_chain,
        reduce_documents_chain=reduce_documents_chain
    )
    
    # Run the MapReduce Chain
    summarization_results = map_reduce_chain.run(documents)

    from langchain.chains.summarize import load_summarize_chain

    llm = ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo")

    chain = load_summarize_chain(llm, chain_type="map_reduce")

    chain.run(documents)

Tokenization

Open-Source alternatives to OpenAI Models

    from langchain.chat_models import ChatMistralAI

    # Initialize LLM
    llm = ChatMistralAI(temperature=0, model_name="mistral-tiny")

    import ollama.client as client

    response, _ = client.generate(model_name="mistral-openorca:latest, system=SYS_PROMPT, prompt=prompt)
    result = json.loads(response)

Other ideas: Knowledge Graph expansion