Vポイントマーケティング|TECH LABの Tech Blog

TECH LABのエンジニアが技術情報を発信しています

ブログタイトル

GraphRAGで生成したKnowledge-GraphをNetworkXとPlotlyで描画しました!

こんにちは、CCCMKホールディングス TECH LAB三浦です。

秋の晴れの日は本当に気持ちがいいですね!特に用もないのにふらっと外を歩きたくなります。

はじめに

さて前回の記事ではMicrosoft Researchの"GraphRAG"というPythonのライブラリをご紹介しました。GraphRAGを使うと、LLMを利用してテキストデータからテキスト内に登場するEntity(物事や出来事、概念など)を抜き出し、2つのEntityの間の関係性を表すKnowledge-Graphを生成することが出来ました。

techblog.cccmkhd.co.jp

ところが前回ではKnowledge-Graphの描画で日本語が文字化けしていたり、ラベルがごちゃごちゃしていて見にくかったりとあまり上手く出来ていませんでした。今回はその課題に取り組んでみました。

今回作ったグラフ

今回作ったグラフです。

ブログから生成したKnowledge-Graphです。

Plotlyを使っているため、特定のエリアをズームしたり、マウスカーソルを合わせるとテキストが表示されるような、インタラクティブ性も持っています。

CCCMKホールディングスTECH LABと三浦の関係性

前回は表示出来ていなかった日本語も、文字化けすることなく表示することが出来ています

このグラフを描くために作ったコードを、ご紹介したいと思います。

描画に使うデータ

最初に描画に使用するデータについて説明します。

前回の記事でご紹介した手順でGraphRAGでIndexingの処理を走らせると、create_final_relationships.parquetというpandas.Dataframeのデータが格納されたファイルが生成されます。

データは以下のような形式になっています。

`create_final_relationships.parquet`

この中からsource, target, description, weightというカラムをグラフ描画用のデータとして利用します。sourcetargetは関係性の始点と終点のEntityを、descriptionは関係性の説明、weightは重要性を表しています。

作業の流れ

Knowledge-Graphを描くためにNetworkXPlotlyというライブラリを使用します。

まずNetworkXを使ってsourcetargetの情報からレイアウトアルゴリズムによって適切な位置にGraphのNodeとして各Entityを配置し、2次元の座標情報を取得します。そしてPlotlyで散布図を描く要領でNodeを描画、線グラフを描く要領でEdgeを描画、そしてEdgeで結ばれた2点のNodeのちょうど真ん中にdescriptionをマウスカーソルが近づいた時だけ表示されるように描画します。

実際のコード

ここからは描画のために作成したコードを掲載します。

ライブラリのインストール

pipコマンドでnetworkxplotlyをインストールします。plotlyはバージョンが古いと有向グラフを表現する矢印付きのEdgeが描けませんでした。5.24.1であれば問題なく動くと思います。

pip install networkx plotly==5.24.1

NetworkXによるNodeの配置

最初にNetworkXを使ってGraphのNodeの座標を生成します。

import networkx as nx
import pandas as pd

relations = pd.read_parquet(
    "./path/to/create_final_relationships.parquet"
)

# 有向グラフ
G = nx.DiGraph()

for source, target in zip(
        relations['source'], 
        relations['target'], 
    ):
    G.add_edge(source, target)

# 座標の計算
pos = nx.fruchterman_reingold_layout(G,k=0.5)

posの中身は次のようになっています。

Nodeの名前とx座標とy座標が格納されています。

使いやすいように次のような関数に座標取得処理をまとめておきました。

def get_node_pos(node_name):
    """
    Nodeの名前からNodeの座標を取得する関数
    node_name: Nodeの名前
    """
    return (pos[node_name][0], pos[node_name][1])

EdgeとTextの描画設定

次にGraphのNode間をつなぐEdgeと、Node間の関係性を表すTextの描画設定を行いました。 Edgeはその重要性を太さで表現するようにし、Textは常に表示されているとグラフがごちゃごちゃして分かりにくいので、常に透過表示しておき、マウスカーソルが近づいた時だけホバー表示するようにしました。

それぞれ描画設定用の処理を関数にまとめてみました。

import plotly.graph_objects as go

def create_edge_trace(x0, y0, x1, y1, weight):
    """
    矢印付のEdgeを描画するための設定関数
    x0, y0: 始点の座標
    x1, y1: 終点の座標
    weight: 関係性の重要性(エッジの太さ)
    """
    edge_trace = go.Scatter(
        x = [x0, x1],
        y = [y0, y1],
        mode = "lines+markers",
        hoverinfo='none',
        # weightの値をそのまま使うとEdgeが太くなりすぎるので0.2倍して調整
        line=dict(width=weight*0.2), 
        marker=dict(
            # 矢印マークの設定
            symbol="arrow-bar-up",
            angleref="previous"
        )
    )
    return edge_trace

def create_text_trace(x0, y0, x1, y1, desc):
    """
    Node間の関係性を示すテキストをホバー表示するための設定関数
    x0, y0: 始点の座標
    x1, y1: 終点の座標
    desc: 関係性を表すテキスト
    """
    text_trace = go.Scatter(
        #2Nodeのちょうど真ん中に表示されるようにする。
        x=[(x0 + x1) / 2],
        y=[(y0 + y1) / 2],
        mode='text',
        text=[desc],
        textposition='middle center',
        hoverinfo='text',
        # opacityを0にすることで透明になる。
        opacity=0,
        textfont=dict(size=7)
    )
    return text_trace

この関数を利用する全体の処理は次のようになります。

edge_traces = []
text_traces = []

for source, target, desc, weight in zip(
        relations['source'], 
        relations['target'], 
        relations['description'],
        relations['weight']
    ):

    # Nodeの名前からそれぞれの座標を取得
    x0, y0 = get_node_pos(source)
    x1, y1 = get_node_pos(target)

    # 有向グラフの矢印付のエッジ設定
    edge_trace = create_edge_trace(x0, y0, x1, y1, weight)

    # 関係性テキストの設定
    text_trace = create_text_trace(x0, y0, x1, y1, desc)
    
    edge_traces.append(edge_trace)
    text_traces.append(text_trace)

Nodeの描画設定

次はNodeの描画設定です。こちらも描画の設定を行う関数を作成しました。

def create_node_trace(x, y, node_name):
    """
    Nodeの描画設定をする関数
    x, y: Nodeの座標
    node_name: Nodeの名前
    """
    node_trace = go.Scatter(
        x=[x],
        y=[y],
        text=node_name,
        mode='markers+text',
        textposition="bottom center",
        marker=dict(
            color='LightSkyBlue',
            size=3,
            line=dict(
                color='MediumPurple',
                width=2
            )
        ),
        textfont=dict(size=10),
        showlegend=False 
    )
    return node_trace

この関数を利用して、全てのNodeの描画設定を行いました。

node_traces = []

for node_name in G.nodes:
    x, y = get_node_pos(node_name)
    node_trace = create_node_trace(x, y, node_name)
    node_traces.append(node_trace)

Graphの描画

最後にGraph全体のレイアウトを設定して、Graphを表示します。

layout = go.Layout(
    title="ブログ⇒Knowledge-Graph",
    titlefont_size=16,
    showlegend=False,
    xaxis_visible=False,
    yaxis_visible=False,
    width=900,
    height=900
)

go.Figure(
    data = edge_traces + node_traces + text_traces,
    layout=layout
)

Graphが描画されました。

Graphを描画することが出来ました。

まとめ

今回はKnowledge-GraphをNetworkXとPlotlyで描画する方法について調べたのでまとめてみました。GraphRAGで生成された結果はなかなか面白い内容になっているのになかなか上手くそれを表現できずに試行錯誤していたのですが、なんとかいい感じに表現できるようになって良かったと思っています。

あとはこのKnowledge-Graphを使って如何に良い回答が出来るLLMアプリケーションを作ることが出来るのか、といったテーマに本格的にトライしていきたいと思います!