DatabricksでNebulaGraphを用いた知識グラフベースRAGを動かしてみる(日本語編)

こちらの続きです。

https://qiita.com/taka_yayoi/items/f5e26b06af038082a4b6

日本語の文書を使って知識グラフベースのRAGを動かします。

[前回と同様](https://qiita.com/taka_yayoi/items/f5e26b06af038082a4b6#nebulagraph%E3%81%AE%E3%82%A4%E3%83%B3%E3%82%B9%E3%83%88%E3%83%BC%E3%83%AB)にDatabricksクラスターにNebulaGraphがインストールされるようにします。

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

```py
%pip install -U llama-index
%pip install ipython-ngql nebula3-python pyvis
%pip install "databricks-sql-connector[sqlalchemy]"
dbutils.library.restartPython()
```

# ライブラリのインポート

```py
import os

# OpenAI APIキー
os.environ["OPENAI_API_KEY"] = dbutils.secrets.get("demo-token-takaaki.yayoi", "openai_api_key")

import logging
import sys

logging.basicConfig(
    stream=sys.stdout, level=logging.DEBUG #logging.INFO
)  # より詳細なログを出力するには logging.DEBUG

from llama_index import (
    KnowledgeGraphIndex,
    ServiceContext,
    SimpleDirectoryReader,
)
from llama_index.storage.storage_context import StorageContext
from llama_index.graph_stores import NebulaGraphStore
from llama_index.llms import OpenAI

from IPython.display import Markdown, display

# LLMを定義
llm = OpenAI(temperature=0, model="gpt-3.5-turbo-instruct")
service_context = ServiceContext.from_defaults(llm=llm, chunk_size_limit=512)
```

# NeburaGraphとの接続

```py
os.environ["NEBULA_USER"] = "root"
os.environ["NEBULA_PASSWORD"] = "nebula"  # デフォルトは "nebula"
os.environ[
    "NEBULA_ADDRESS"
] = "127.0.0.1:9669"  # ローカルにNeburaGraphがインストールされている前提

space_name = "llamaindex"
edge_types, rel_prop_names = ["relationship"], [
    "relationship"
]  # デフォルト値、空の知識グラフから作成する際には省略可能
tags = ["entity"]  # デフォルト値、空の知識グラフから作成する際には省略可能
```

## スペースの作成

```py
%load_ext ngql
%ngql --address 127.0.0.1 --port 9669 --user root --password nebula
```

```py
%ngql add hosts 127.0.0.1:9779
```

すぐに実行するとエラーになるのでスリープ時間を挿入します。

```py
import time
time.sleep(10)
```

```py
%ngql CREATE SPACE llamaindex(vid_type=FIXED_STRING(256), partition_num=1, replica_factor=1);
```

```py
time.sleep(10)
```

## タグとエッジ、インデックスの作成

```py
%ngql USE llamaindex;
%ngql CREATE TAG entity(name string);
%ngql CREATE EDGE relationship(relationship string);
```

```py
time.sleep(10)
```

```py
%ngql CREATE TAG INDEX entity_index ON entity(name(256));
```

## グラフストアの作成

```py
graph_store = NebulaGraphStore(
    space_name=space_name,
    edge_types=edge_types,
    rel_prop_names=rel_prop_names,
    tags=tags,
)
storage_context = StorageContext.from_defaults(graph_store=graph_store)
```

# Unity Catalogとの接続

```py
from llama_index import download_loader
from llama_index.utilities.sql_wrapper import SQLDatabase

DatabaseReader = download_loader('DatabaseReader')

# パーソナルアクセストークン
access_token    = dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiToken().getOrElse(None)
# Databricksワークスペースホスト名
server_hostname = "<ワークスペースのホスト名>"
# クラスターのHTTPパス
http_path       = "sql/protocolv1/o/2556758628403379/1118-004519-nu0u899x"
# Unity Catalogカタログ
catalog         = "takaakiyayoi_catalog"
# スキーマ(データベース)
schema          = "qiita_2023"

reader = DatabaseReader(
    sql_database = SQLDatabase.from_uri(f"databricks://token:{access_token}@{server_hostname}?" +
        f"http_path={http_path}&catalog={catalog}&schema={schema}")
)
```

# テキストデータのロード

```py
# テキストを格納しているカラムに応じてSQLを変更します
query = f"""
SELECT
    body
FROM takaakiyayoi_catalog.qiita_2023.taka_qiita_2023 LIMIT 100
"""

documents = reader.load_data(query=query)
```

# 知識グラフインデックスの作成

それなりに時間を要するので、最初はドキュメントの数を減らした方がいいかもしれません。

```py
kg_index = KnowledgeGraphIndex.from_documents(
    documents,
    storage_context=storage_context,
    max_triplets_per_chunk=10,
    service_context=service_context,
    space_name=space_name,
    edge_types=edge_types,
    rel_prop_names=rel_prop_names,
    tags=tags,
    include_embeddings=True,
)
```

# 知識グラフへのクエリー

```py
%ngql USE llamaindex;
%ngql MATCH ()-[e]->() RETURN e LIMIT 10
```
![Screenshot 2024-01-23 at 11.34.11.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/daf641f0-c7e0-142a-38e1-41f807f48218.png)

```py
# draw the result
%ng_draw
```

前回はノートブックにレンダリングできませんでしたが、以下の方法でレンダリングできました。

```py
html_file_content = open("/Workspace/Users/takaaki.yayoi@databricks.com/20240122_nebula_llama/nebulagraph.html", 'r').read()
displayHTML(html_file_content)
```
![Screenshot 2024-01-23 at 11.43.37.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/2596c114-8d67-9885-90a0-3b29c2cbcc14.png)

# RAGへの問い合わせ

```py
from llama_index.query_engine import KnowledgeGraphQueryEngine

from llama_index.storage.storage_context import StorageContext
from llama_index.graph_stores import NebulaGraphStore

query_engine = KnowledgeGraphQueryEngine(
    storage_context=storage_context,
    service_context=service_context,
    llm=llm,
    verbose=True,
)
```

```py
response = query_engine.query(
    "Databricksとは",
)
display(Markdown(f"<b>{response}</b>"))
```
![Screenshot 2024-01-23 at 11.44.44.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/a7d7c1f0-8128-d688-9af8-5af4214c4942.png)


```py
response = query_engine.query(
    "Apache Sparkとは",
)
display(Markdown(f"<b>{response}</b>"))
```

動いてますね。
![Screenshot 2024-01-23 at 11.45.09.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/ea7b7c78-6099-0a0d-dd6f-e6d793d885ef.png)

もう少し挙動を理解してチューニングできるところあればしたいところです。

### はじめてのDatabricks

[はじめてのDatabricks](https://qiita.com/taka_yayoi/items/8dc72d083edb879a5e5d)

### Databricks無料トライアル

[Databricks無料トライアル](https://databricks.com/jp/try-databricks)