사전 정의된 스키마(Pre-defined Schemas)로 Property Graph 구축하기

Neo4j, Ollama, Huggingface를 사용해 프로퍼티 그래프를 구축하는 문서예요. SchemaLLMPathExtractor로 가능한 엔터티 타입, 관계 타입, 그리고 그 연결 방식을 담은 정밀한 스키마를 정의해, LLM이 임의로 결정하는 대신 사전 정의된 스키마에 맞는 엔터티·관계를 추출해요.

출처: 문서

본문

이 노트북에서는 Neo4j, Ollama, Huggingface를 사용해 프로퍼티 그래프를 구축하는 과정을 안내해요. 특히 SchemaLLMPathExtractor를 활용해, 가능한 엔터티 타입과 관계 타입, 그리고 그들이 어떻게 연결될 수 있는지가 담긴 정밀한 스키마를 정의해요. LLM이 엔터티와 관계를 임의로 정하게 두는 대신, 이 사전 정의된 스키마에 부합하는 엔터티와 관계만 추출하게 해요.

설치 (Setup)

%pip install llama-index-core
%pip install llama-index-llms-ollama
%pip install llama-index-embeddings-huggingface
%pip install llama-index-graph-stores-neo4j

import nest_asyncio
nest_asyncio.apply()
from IPython.display import Markdown, display

데이터 다운로드/로드 (Download Data / Load Data)

!mkdir -p 'data/paul_graham/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'

from llama_index.core import SimpleDirectoryReader
documents = SimpleDirectoryReader("./data/paul_graham/").load_data()

그래프 생성 (Graph Construction)

그래프를 만들려면 SchemaLLMPathExtractor를 사용해요. 이 도구로 그래프의 스키마를 지정할 수 있고, LLM이 엔터티와 관계를 임의로 정하게 두는 대신 사전 정의된 스키마에 맞는 엔터티·관계를 추출하게 해줘요.

from typing import Literal
from llama_index.llms.ollama import Ollama
from llama_index.core.indices.property_graph import SchemaLLMPathExtractor

# best practice to use upper-case
entities = Literal["PERSON", "PLACE", "ORGANIZATION"]
relations = Literal["HAS", "PART_OF", "WORKED_ON", "WORKED_WITH", "WORKED_AT"]

# define which entities can have which relations
# validation_schema = {
#     "PERSON": ["HAS", "PART_OF", "WORKED_ON", "WORKED_WITH", "WORKED_AT"],
#     "PLACE": ["HAS", "PART_OF", "WORKED_AT"],
#     "ORGANIZATION": ["HAS", "PART_OF", "WORKED_WITH"],
# }
validation_schema = [
    ("ORGANIZATION", "HAS", "PERSON"),
    ("PERSON", "WORKED_AT", "ORGANIZATION"),
    ("PERSON", "WORKED_WITH", "PERSON"),
    ("PERSON", "WORKED_ON", "ORGANIZATION"),
    ("PERSON", "PART_OF", "ORGANIZATION"),
    ("ORGANIZATION", "PART_OF", "ORGANIZATION"),
    ("PERSON", "WORKED_AT", "PLACE"),
]

kg_extractor = SchemaLLMPathExtractor(
    llm=Ollama(model="mistral:7b", json_mode=True, request_timeout=3600),
    possible_entities=entities,
    possible_relations=relations,
    kg_validation_schema=validation_schema,
    # if false, allows for values outside of the schema
    # useful for using the schema as a suggestion
    strict=True,
)

Neo4j 설정 (Neo4j setup)

Docker로 Neo4j 컨테이너를 실행해요.

!docker run \
    -p 7474:7474 -p 7687:7687 \
    -v $PWD/data:/data -v $PWD/plugins:/plugins \
    --name neo4j-apoc \
    -e NEO4J_apoc_export_file_enabled=true \
    -e NEO4J_apoc_import_file_enabled=true \
    -e NEO4J_apoc_import_file_use__neo4j__config=true \
    -e NEO4JLABS_PLUGINS=\[\"apoc\"\] \
    neo4j:latest

PropertyGraphIndex 생성 (PropertyGraphIndex construction)

from llama_index.graph_stores.neo4j import Neo4jPropertyGraphStore
from llama_index.core.vector_stores.simple import SimpleVectorStore

graph_store = Neo4jPropertyGraphStore(
    username="neo4j",
    password="llamaindex",
    url="bolt://localhost:7687",
)
vec_store = SimpleVectorStore()

from llama_index.core import PropertyGraphIndex
from llama_index.embeddings.huggingface import HuggingFaceEmbedding

index = PropertyGraphIndex.from_documents(
    documents,
    kg_extractors=[kg_extractor],
    embed_model=HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5"),
    property_graph_store=graph_store,
    vector_store=vec_store,
    show_progress=True,
)

리트리버 설정 (Setup Retrievers)

동의어 리트리버와 벡터 컨텍스트 리트리버를 함께 사용해요.

from llama_index.core.indices.property_graph import (
    LLMSynonymRetriever,
    VectorContextRetriever,
)

llm_synonym = LLMSynonymRetriever(
    index.property_graph_store,
    llm=Ollama(model="mistral:instruct", request_timeout=3600),
    include_text=False,
)
vector_context = VectorContextRetriever(
    index.property_graph_store,
    embed_model=HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5"),
    include_text=False,
)

retriever = index.as_retriever(
    sub_retrievers=[
        llm_synonym,
        vector_context,
    ]
)

질의 (Querying)

리트리버 (Retriever)

nodes = retriever.retrieve("What happened at Interleaf?")
for node in nodes:
    print(node.text)

쿼리 엔진 (QueryEngine)

query_engine = index.as_query_engine(
    sub_retrievers=[
        llm_synonym,
        vector_context,
    ],
    llm=Ollama(model="mistral:instruct", request_timeout=3600),
)
response = query_engine.query("What happened at Interleaf?")
display(Markdown(f"{response.response}"))

이처럼 사전 정의된 스키마로 그래프 구축을 제약하면 더 깔끔하고 예측 가능한 지식 그래프를 만들 수 있어요. strict=True로 두면 스키마 밖의 값이 허용되지 않아요.

사전 정의된 스키마(Pre-defined Schemas)로 Property Graph 구축하기