Tiny GraphRAG (2부)
1부에서는 GraphRAG의 Core 개념을 보여 주는 최소한의 구현을 만들었습니다. 이제 프로덕션 환경에 더 적합하도록 세 가지 중요한 개선 사항으로 구현을 확장하겠습니다.
- 직렬화된 networkx 그래프를 본격적인 그래프 데이터베이스인 Memgraph로 교체합니다.
- 문서의 자연스러운 구조를 더 잘 포착하기 위해 계층적 커뮤니티 탐지를 추가합니다.
- 동일한 엔터티에 대한 서로 다른 언급을 표준 형태로 귀결시키기 위해 엔터티 명확화를 구현합니다.
이러한 기능의 전체 구현은 저장소의 micro-graphrag 브랜치에서 확인할 수 있습니다. 새로 추가된 내용을 포함하면 확장된 구현은 약 1500줄의 코드에 이릅니다.
Memgraph 구성
Memgraph를 그래프 데이터베이스로 사용하기 위해 docker-compose.yml에 필요한 서비스를 추가하겠습니다.
services:
memgraph:
image: memgraph/memgraph-mage:1.19-memgraph-2.19
container_name: memgraph
ports:
- "7687:7687" # Bolt protocol port
- "7444:7444" # HTTP API port
command: ["--log-level=TRACE", "--query-execution-timeout-sec=0"]
volumes:
- mg_lib:/var/lib/memgraph
lab:
image: memgraph/lab:latest
container_name: memgraph-ui
ports:
- "3000:3000"
depends_on:
- memgraph
environment:
- QUICK_CONNECT_MG_HOST=memgraph
- QUICK_CONNECT_MG_PORT=7687
- QUICK_CONNECT_MG_USER=admin
- QUICK_CONNECT_MG_PASSWORD=admin
volumes:
mg_lib:
이 구성은 다음과 같은 일을 합니다.
- 고급 그래프 알고리즘을 위한 MAGE 확장과 함께 Memgraph 데이터베이스를 설정합니다
- 그래프를 시각화하고 쿼리하기 위한 웹 기반 UI(Memgraph Lab)를 제공합니다
- Docker 볼륨을 통해 영구 저장소를 구성합니다
- 데이터베이스(7687)와 UI(3000) 모두에 필요한 포트를 노출합니다
연결 세부 정보는 GraphConfig 클래스에서 구성합니다.
@dataclass
class GraphConfig:
"""Memgraph connection configuration."""
uri: str = "bolt://localhost:7687"
username: str = "admin"
password: str = "admin"
database: str = "memgraph"
이 설정은 기존의 PostgreSQL/pgvector 서비스를 보완하여 다음을 제공합니다.
- 벡터 임베딩과 검색을 위한 pgvector 기반 PostgreSQL
- 효율적인 그래프 저장과 쿼리를 위한 Memgraph
- 그래프 시각화와 탐색을 위한 웹 UI
이 조합을 통해 GraphRAG 구현에서 벡터 연산과 그래프 연산을 모두 효과적으로 활용할 수 있습니다.
빌드 파이프라인 변경 사항(build.py)
빌드 파이프라인은 향상된 기능을 지원하기 위해 새로운 데이터클래스와 메서드로 업데이트되었습니다. Core 자료 구조는 문서 처리 파이프라인을 더 잘 표현할 수 있도록 확장되었습니다.
@dataclass
class Entity:
"""Graph node with type information."""
id: str
type: str
@dataclass
class Relation:
"""Graph triple representing relationships."""
head: str
relation_type: str
tail: str
@dataclass
class DocumentChunkData:
"""Document chunk with embedding."""
text: str
embedding: Any
@dataclass
class ProcessedDocument:
"""Complete processed document data."""
chunks: List[DocumentChunkData]
entities: List[Entity]
relations: List[Relation]
주요 문서 처리 파이프라인은 엔터티 명확화와 개선된 그래프 구성을 지원하도록 크게 향상되었습니다. 이제 이 과정은 두 번의 패스로 진행됩니다.
- 첫 번째 패스에서는 모든 엔터티 언급을 해당 문맥과 함께 수집합니다
- 두 번째 패스에서는 확정된 표준 엔터티를 사용하여 그래프를 구축합니다
다음은 Core 구현입니다:
def process_document(
filepath: str,
title: Optional[str] = None,
max_chunks: int = -1,
entity_types: List[str] = MIN_ENTITY_TYPES,
relation_types: List[str] = DEFAULT_RELS_LIST,
) -> ProcessedDocument:
"""Process document with entity disambiguation."""
# First pass: collect all entity mentions
all_mentions = []
for chunk_text, _ in page_chunks[:max_chunks]:
extraction = extract_rels(chunk_text, entity_types, relation_types)
for ent in extraction.entities:
all_mentions.append((ent[0], ent[1], chunk_text))
# Resolve entities to canonical forms
disambiguator = EntityDisambiguator(model)
resolved_entities = disambiguator.resolve_entities(all_mentions)
# Map mentions to canonical forms
entity_map = {
m[0]: r[0]
for m, r in zip(all_mentions, resolved_entities)
}
# Second pass: create graph with resolved entities
g = nx.Graph()
for chunk_text, _ in page_chunks[:max_chunks]:
extraction = extract_rels(chunk_text, entity_types, relation_types)
for rel in extraction.relations:
if rel[0] in entity_map and rel[2] in entity_map:
head = entity_map[rel[0]]
tail = entity_map[rel[2]]
g.add_edge(head, tail,
label=rel[1],
source_chunk=chunk_text
)
return g
두 패스 방식을 사용하면 명확화 결정을 내리기 전에 가능한 모든 엔터티 언급을 문맥과 함께 수집하고, 표준 엔터티 형태만 사용하여 깔끔한 그래프 구조를 만들며, 명확화된 엔터티 간의 관계 매핑을 올바르게 유지할 수 있습니다.
이 프로세스는 GraphStore 클래스를 통해 새로운 Memgraph 저장 시스템과 통합되며, 이 클래스가 그래프 데이터의 저장과 조회를 담당합니다. 이렇게 개선된 파이프라인은 문맥을 고려한 중의성 해소로 엔터티 해석을 개선합니다. 또한 중복이 제거된 엔터티로 더 깔끔한 그래프 구조를 만들고, 정규 엔터티 형태로 관계의 정확도를 높이며, Memgraph 통합으로 효율적인 저장과 조회를 제공합니다.
이제 빌드 과정은 엔터티의 변형과 중의성을 처리하면서도 문서의 의미적 관계를 더 잘 포착하는, 더욱 견고한 지식 그래프를 만들어 냅니다. 이렇게 개선된 그래프 구조는 더 정확한 엔터티 관계와 더 잘 정리된 커뮤니티 구조를 제공하므로 로컬 검색과 글로벌 검색 성능을 모두 직접적으로 향상시킵니다.
Memgraph 통합 (graph.py)
원래 구현에서는 직렬화된 networkx 그래프를 pickle 파일로 저장하여 사용했습니다. 이 방식은 시연 용도로는 충분했지만 프로덕션 환경에는 적합하지 않습니다. 이를 머신러닝 워크로드에 최적화된 고성능 그래프 데이터베이스인 Memgraph로 교체합시다.
pickle 기반 저장소를 Memgraph로 모든 그래프 연산을 처리하는 새로운 GraphStore 클래스로 교체하겠습니다. 이 클래스는 그래프 데이터를 저장하고 조회하기 위한 깔끔한 인터페이스를 제공합니다. 엔터티와 관계를 추가하고, 그래프 구조를 쿼리하고, 문서별 하위 그래프를 관리하는 메서드도 갖추고 있습니다. Memgraph 백엔드는 최적화된 그래프 순회 알고리즘과 병렬 쿼리 기본 지원 덕분에 성능을 크게 향상시켜 줍니다. 다음은 Core 구현입니다:
class GraphStore:
"""Handles graph storage and retrieval using Memgraph."""
def __init__(self, config: GraphConfig | None = None):
self.driver = GraphDatabase.driver(
(config or GraphConfig()).uri,
auth=((config or GraphConfig()).username,
(config or GraphConfig()).password)
)
def store_graph(
self,
doc_id: int,
entities: List[Tuple[str, str]],
relations: List[Tuple[str, str, str]],
source_chunk: str,
) -> None:
"""Store entities and relations in Memgraph."""
with self.driver.session() as session:
# Store entities with vector embeddings
for text, label in entities:
query = """
MERGE (n:Entity {
content: $content,
label: $label,
doc_id: $doc_id
})
"""
session.run(query, {
"content": text,
"label": label,
"doc_id": doc_id
})
# Store relations with source context
for head, rel_type, tail in relations:
query = """
MATCH (h:Entity {content: $head}),
(t:Entity {content: $tail})
CREATE (h)-[:RELATES {
type: $rel_type,
doc_id: $doc_id,
source_chunk: $source_chunk
}]->(t)
"""
session.run(query, {
"head": head,
"tail": tail,
"rel_type": rel_type,
"doc_id": doc_id,
"source_chunk": source_chunk
})
계층적 커뮤니티 (communities.py)
원래의 커뮤니티 탐지 구현은 각 엔터티가 정확히 하나의 커뮤니티에만 속하는 평면 구조를 사용했습니다. 이 방식은 주제들이 자연스럽게 계층적 관계를 이루는 대규모 문서에서는 한계가 될 수 있습니다. 각 엔터티가 정확히 하나의 커뮤니티에만 속하던 이전의 평면 구현과 달리, 이 방식은 여러 수준의 커뮤니티를 만들며 큰 커뮤니티가 더 작은 하위 커뮤니티를 포함합니다.
예를 들어 오바마에 관한 위키백과 문서를 분석하면 다음과 같은 계층 구조를 얻을 수 있습니다:
- 수준 0: 주요 생애 시기
- 정치 경력 커뮤니티
- 수준 1: 대통령 임기
- 첫 번째 임기 커뮤니티
- 수준 2: 주요 사건
- 금융 위기 대응
- 의료 개혁
- 빈 라덴 작전
- 두 번째 임기 커뮤니티
- 수준 2: 주요 사건
- 기후 변화 대응 조치
- 이란 핵 합의
- 쿠바와의 관계
- 수준 1: 대통령 취임 이전 경력
- 상원 경력 커뮤니티
- 주 상원 커뮤니티
- 개인 생활 커뮤니티
- 레벨 1: 교육
- 컬럼비아 대학교
- 하버드 로스쿨
- 레벨 1: 어린 시절
- 하와이 배경
- 시카고 지역사회 활동
HierarchicalCommunity 클래스는 이 트리의 각 커뮤니티 노드를 나타냅니다. 이 클래스는 다음을 저장합니다.
- 해당 커뮤니티에 포함된 트리플(엔티티 관계)
- 바로 아래 레벨의 하위 커뮤니티 목록
- 커뮤니티 내용에 대한 선택적 요약
- 계층 구조에서의 레벨/깊이
build_hierarchical_communities 함수는 다음과 같은 방식으로 이 트리 구조를 재귀적으로 구축합니다.
- 각 서브그래프에서 커뮤니티 탐지를 실행합니다
- 해당 커뮤니티의 관계 트리플을 추출합니다
- 각 커뮤니티 내에서 하위 커뮤니티를 재귀적으로 탐지합니다
- 최대 깊이 또는 최소 커뮤니티 크기에 도달하면 중단합니다
이러한 계층적 접근 방식을 사용하면 문서 구조와 엔티티 간 관계를 다양한 규모에서 더 세밀하게 분석할 수 있습니다. 위의 오바마 예시에서 이 방식이 넓은 생애 시기와 대통령 재임 기간의 구체적인 사건 및 업적을 모두 자연스럽게 포착하는 것을 확인할 수 있습니다.
@dataclass
class HierarchicalCommunity:
"""Represents a hierarchical community structure."""
triples: List[Tuple[str, str, str, str]]
subcommunities: List["HierarchicalCommunity"]
summary: Optional[str] = None
level: int = 0
def build_hierarchical_communities(
g: nx.Graph,
max_levels: int = 3,
min_size: int = 5
) -> CommunityResult:
"""Build hierarchical communities from a graph."""
def build_level(subgraph: nx.Graph, level: int) -> List[HierarchicalCommunity]:
if level >= max_levels or len(subgraph) < min_size:
return []
communities = []
coms = algorithms.leiden(subgraph)
for com in coms.communities:
community_subgraph = subgraph.subgraph(com)
# Extract triples for this community
triples = [
(s, data["label"], t, data["source_chunk"])
for s, t, data in community_subgraph.edges(data=True)
]
# Recursively build subcommunities
subcommunities = build_level(community_subgraph, level + 1)
communities.append(HierarchicalCommunity(
triples=triples,
subcommunities=subcommunities,
level=level
))
return communities
return CommunityResult(
communities=build_level(g, 0),
node_community_map=build_community_map(g)
)
이러한 계층적 접근 방식은 문서 구조를 다양한 세분성 수준에서 다중 규모로 분석할 수 있게 하여 몇 가지 핵심적인 이점을 제공합니다. 주제는 “정치“에서 “대통령 선거 운동“으로, 다시 “예비 선거“로 이어지는 것처럼 의미 있는 계층 구조로 자연스럽게 정리됩니다. 이 접근 방식은 질의를 가장 적절한 세부 수준에 맞춰 연결하므로 유연한 질의가 가능합니다. 또한 문서 구조 안에서 넓은 주제와 구체적인 세부 사항을 모두 포착할 수 있도록 요약 생성을 개선합니다.
엔티티 중의성 해소 (dis.py)
기존 구현에서는 엔티티를 정확한 문자열 일치로만 처리했기 때문에, 같은 엔티티를 가리키는 서로 다른 표현(예: “Barack Obama“와 “President Obama”) 사이의 연결을 놓쳤습니다. 벡터 유사도와 문맥 단서를 활용한 엔티티 중의성 해소를 추가해 봅시다.
다음은 벡터 임베딩과 문맥 정보를 사용하여 같은 엔티티에 대한 서로 다른 언급을 하나의 정규 형태로 통합하는 새로운 엔티티 중의성 해소 시스템입니다.
- 각 엔티티 언급은 sentence-transformers를 사용하여 주변 문맥과 함께 임베딩됩니다.
- 모든 언급 임베딩 간의 유사도 행렬을 계산합니다
- 임계값(0.85)을 넘는 유사도 점수를 기준으로 언급을 클러스터링합니다
- 각 클러스터마다 빈도와 완전성을 기준으로 정규 형태가 선택됩니다
- 엔티티 유형은 클러스터에 나타나는 유형 중 가장 구체적인 유형을 택하는 방식으로 결정됩니다
이 시스템은 다음과 같은 경우를 처리합니다:
- 이름 변형 (“Barack Obama” 대 “President Obama”)
- 직함 및 경칭 (“Dr. Smith” 대 “John Smith”)
- 약어 및 두문자어 (“United Nations” 대 “UN”)
- 문맥적 지칭 (특정 인물을 가리키는 “the president”)
아래 구현은 Core 동음이의어 해소(disambiguation) 로직을 보여 줍니다:
class EntityDisambiguator:
def __init__(self, embedding_model: SentenceTransformer):
self.model = embedding_model
self.mention_cache: Dict[str, np.ndarray] = {}
def get_mention_embedding(self, mention: str, context: str) -> np.ndarray:
"""Get embedding for entity mention with context."""
key = f"{mention}::{context}"
if key not in self.mention_cache:
# Combine mention with surrounding context
text = f"{mention} | {context}"
self.mention_cache[key] = self.model.encode(text)
return self.mention_cache[key]
def resolve_entities(
self,
mentions: List[Tuple[str, str, str]] # (text, type, context)
) -> List[Tuple[str, str]]: # (canonical_form, type)
"""Resolve entity mentions to canonical forms."""
resolved = []
clusters: Dict[str, List[int]] = {}
# Compute similarity matrix
embeddings = [
self.get_mention_embedding(m[0], m[2])
for m in mentions
]
similarities = cosine_similarity(embeddings)
# Cluster similar mentions
for i in range(len(mentions)):
assigned = False
for canonical, cluster in clusters.items():
# Check similarity with existing cluster
mean_sim = np.mean([similarities[i][j] for j in cluster])
if mean_sim > 0.85: # Similarity threshold
cluster.append(i)
assigned = True
break
if not assigned:
# Create new cluster
clusters[mentions[i][0]] = [i]
# Select canonical forms and resolve types
for canonical, cluster in clusters.items():
# Use most frequent entity type in cluster
types = [mentions[i][1] for i in cluster]
resolved.append((canonical, mode(types)))
return resolved
중의성 해소기는 문서 처리 파이프라인에 다음과 같이 통합됩니다:
def process_document(
filepath: str,
title: Optional[str] = None,
max_chunks: int = -1,
entity_types: List[str] = MIN_ENTITY_TYPES,
relation_types: List[str] = DEFAULT_RELS_LIST,
) -> ProcessedDocument:
"""Process document with entity disambiguation."""
# First pass: collect all entity mentions
all_mentions = []
for chunk_text, _ in page_chunks[:max_chunks]:
extraction = extract_rels(chunk_text, entity_types, relation_types)
for ent in extraction.entities:
all_mentions.append((ent[0], ent[1], chunk_text))
# Resolve entities to canonical forms
disambiguator = EntityDisambiguator(model)
resolved_entities = disambiguator.resolve_entities(all_mentions)
# Map mentions to canonical forms
entity_map = {
m[0]: r[0]
for m, r in zip(all_mentions, resolved_entities)
}
# Second pass: create graph with resolved entities
g = nx.Graph()
for chunk_text, _ in page_chunks[:max_chunks]:
extraction = extract_rels(chunk_text, entity_types, relation_types)
for rel in extraction.relations:
if rel[0] in entity_map and rel[2] in entity_map:
head = entity_map[rel[0]]
tail = entity_map[rel[2]]
g.add_edge(head, tail,
label=rel[1],
source_chunk=chunk_text
)
return g
결론
이러한 개선은 시스템의 기능을 크게 향상시키지만, 상당한 계산 비용이 따르는 것도 사실입니다. Memgraph를 사용하면 직렬화된 그래프에 비해 더 많은 메모리가 필요하지만, 이러한 트레이드오프 덕분에 확장성이 더 좋아집니다. 계층적 커뮤니티 탐지 과정은 단순한 평면 클러스터링 방식보다 계산량이 더 많습니다. 또한 엔티티 해석에 필요한 벡터 유사도 비교는 문서 수집 과정에서 추가적인 처리 오버헤드를 발생시킵니다.
이러한 개선이 제공하는 이점은 대체로 추가적인 자원 사용을 정당화합니다. 쿼리 성능이 극적으로 향상되며, Memgraph 덕분에 그래프 쿼리가 이전 방식보다 10배에서 100배까지 빨라집니다. 계층적 커뮤니티 구조와 더 정확한 엔티티 중의성 해소가 결합되어 답변의 품질이 크게 향상됩니다. 무엇보다 중요한 점은, 개선된 아키텍처 덕분에 시스템이 이전보다 훨씬 더 큰 문서와 더 복잡한 지식 그래프를 처리할 수 있도록 효과적으로 확장된다는 것입니다.
이러한 개선으로 우리 구현이 프로덕션 환경에 더 적합해지기는 했지만, 원래의 GraphRAG 논문과 마찬가지로 내부에서는 여전히 “그래프를 넣으면 마법이 나온다“는 식의 일이 꽤 많이 일어나고 있다는 점을 기억해 둘 필요가 있습니다. Memgraph로 영속성을 개선하고, 커뮤니티 탐지를 더 똑똑하게 만들고, 엔티티 처리를 더 견고하게 했지만, Core만 놓고 보면 여전히 본질적으로 여러 언어 모델을 그래프에 던져 넣고 서로 잘 어울리기를 바라는 것에 불과합니다. 때로는 잘 협력하고, 때로는 그렇지 않으며, 때로는 “우연한 유능함“이라고밖에 표현할 수 없는 방식으로 놀라울 만큼 좋은 결과를 내기도 합니다.