🚀 LightRAG: シンプルかつ高速なRAGフレームワーク(GoogleColab付)

LightRAG公式ページ | arXiv:2410.05779 | LearnOpenCVでの紹介

LightRAGは、テキストやナレッジグラフ、ベクターストアを活用して効率的なRAGワークフローを可能にするフレームワークです。
ここでは、Google Colab風にセルごとに解説とコードを整理し、日本語で手軽にLightRAGを始められる手順を紹介します。

ポイント:

ドキュメント、PDF、CSV、Docxなど様々な形式のコンテンツを取り込み可能
OpenAI API互換、HuggingFaceモデル、Ollamaモデルなど多様なLLMに対応
取得したコンテキストを活用し、高度な質問に自然言語で回答
Neo4Jなどの商用付加ストレージは使用せず、デフォルトのJsonなどを利用可能（簡素化のためNeo4J例は省略します）
グラフ可視化なども可能ですが、ここでは別途可視化コードを用意することにします

インストールと準備

まずはLightRAGをインストールします。
PyPIからもインストール可能ですが、ここではソースからのインストールを推奨します。

LightRAGを取り敢えずGoogle colabで動かしてみます③
取り敢えず4000行の文章を2000行にしたらいけました！！
普通にRAGの挙動見てるけど良さそう！！ https://t.co/wnA33dPUvT pic.twitter.com/xC9bpwHiYB

— Maki@Sunwood AI Labs. (@hAru_mAki_ch) December 10, 2024

インストール(ソースから)

# LightRAGのリポジトリをcloneしてインストール（例）
# ※本セルはGoogle Colab上で実行することを想定しています。
# もしGitが使えない環境やリポジトリURLが必要なら下記コメントを外して適宜調整してください。
!git clone https://github.com/HKUDS/LightRAG.git

%cd LightRAG
!pip install -e .
# ↑ ソースからインストールの例

# ここではPyPIからインストールする例を示します:
# pip install lightrag-hku

import os
from google.colab import userdata

os.environ["OPENAI_API_KEY"] = userdata.get('OPENAI_API_KEY')

LightRAGの基本的な使い方

LightRAGは、作業用ディレクトリを指定し、LLMやEmbedding関数を設定することで動作します。
以下は、ダミーテキスト（例：チャールズ・ディケンズの「クリスマスキャロル」）を取り込み、RAG問い合わせを行う基本例です。

※OpenAIモデルを使用する場合、OPENAI_API_KEY の環境変数が必要です。

データの準備

# デモ用のテキストファイルを取得（英語の小説）
!curl https://raw.githubusercontent.com/gusye1234/nano-graphrag/main/tests/mock_data.txt > ./book.txt

基本的なセットアップとデータ挿入

import nest_asyncio
nest_asyncio.apply()

import os
from lightrag import LightRAG, QueryParam
from lightrag.llm import gpt_4o_mini_complete

WORKING_DIR = "./dickens"
if not os.path.exists(WORKING_DIR):
    os.mkdir(WORKING_DIR)

# LLMモデル関数にはデフォルトでgpt_4o_mini_completeを使います
rag = LightRAG(
    working_dir=WORKING_DIR,
    llm_model_func=gpt_4o_mini_complete
)

# テキストの挿入
with open("./book.txt") as f:
    rag.insert(f.read())

解説:
LightRAGインスタンスを作成し、rag.insert()でコンテンツをベクターストアやKVストア、グラフに登録します。
これで後のクエリ時に、このテキストの知識を活用可能になります。

クエリの実行

rag.query()メソッドを使って、モデルに質問を投げかけます。
QueryParamを指定することで、local, global, hybrid, naive などの検索モードを選べます。

検索モードと質問の例

# テキストに関するテーマを尋ねる例
query_text = "この物語の主要なテーマは何ですか？"

# naive検索（単純検索）
print("===Naive検索===")
print(rag.query(query_text, param=QueryParam(mode="naive")))

# local検索（ローカルエンティティ重視）
print("===Local検索===")
print(rag.query(query_text, param=QueryParam(mode="local")))

# global検索（グローバル関係重視）
print("===Global検索===")
print(rag.query(query_text, param=QueryParam(mode="global")))

# hybrid検索（ローカル+グローバル混合）
print("===Hybrid検索===")
print(rag.query(query_text, param=QueryParam(mode="hybrid")))

解説:
modeによって、RAGが返すコンテキストや検索方法が変化します。

naive：単純な類似検索
local：ローカルコンテキスト（エンティティ）重視
global：グローバル関係性重視
hybrid：ローカル・グローバル両面から考慮

OpenAIライクなAPIの使用例（任意）

OpenAI互換API形式のLLMやEmbeddingsを使用したい場合、下記のような関数を設定できます。
（ここではサンプルコード紹介のみ、実際に動かす場合は適宜API Keyが必要）

OpenAI互換形式の使用例（解説のみ）

# こちらは実行例ではなくサンプルコードです。
# APIキーなどは環境変数で設定してください。
# from lightrag.llm import openai_complete_if_cache, openai_embedding
# import numpy as np
#
# async def llm_model_func(
#     prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
# ) -> str:
#     return await openai_complete_if_cache(
#         "solar-mini",
#         prompt,
#         system_prompt=system_prompt,
#         history_messages=history_messages,
#         api_key=os.getenv("UPSTAGE_API_KEY"),
#         base_url="https://api.upstage.ai/v1/solar",
#         **kwargs
#     )
#
# async def embedding_func(texts: list[str]) -> np.ndarray:
#     return await openai_embedding(
#         texts,
#         model="solar-embedding-1-large-query",
#         api_key=os.getenv("UPSTAGE_API_KEY"),
#         base_url="https://api.upstage.ai/v1/solar"
#     )
#
# # 上記関数を使用するには、ragインスタンス作成時に指定
# # rag = LightRAG(
# #     working_dir=WORKING_DIR,
# #     llm_model_func=llm_model_func,
# #     embedding_func=EmbeddingFunc(
# #         embedding_dim=4096,
# #         max_token_size=8192,
# #         func=embedding_func
# #     )
# # )

Hugging FaceモデルやOllamaモデルの利用

Hugging FaceモデルやOllamaモデルでも同様の設定が可能です。
GPU環境でローカルLLMを使用したい場合、Hugging FaceモデルやOllama対応が便利です。

Hugging Faceモデル利用例（解説のみ）

# 実行例：HuggingFaceモデルを用いる場合
# from transformers import AutoModel, AutoTokenizer
# from lightrag.llm import hf_model_complete, hf_embedding
# from lightrag.utils import EmbeddingFunc
#
# rag_hf = LightRAG(
#     working_dir="./hf_demo",
#     llm_model_func=hf_model_complete,  # Hugging Faceモデルでのテキスト生成
#     llm_model_name='meta-llama/Llama-3.1-8B-Instruct',  # HFのモデル名
#     embedding_func=EmbeddingFunc(
#         embedding_dim=384,
#         max_token_size=5000,
#         func=lambda texts: hf_embedding(
#             texts,
#             tokenizer=AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2"),
#             embed_model=AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
#         )
#     ),
# )

複数ファイルの取り込みやテキスト以外のファイルのサポート

textractを使ってPDFやDocxなどをテキスト化して挿入可能です。
下記はコード例だけ示します。

Multi-file対応（解説のみ）

# import textract
# file_path = 'TEXT.pdf'
# text_content = textract.process(file_path)
# rag.insert(text_content.decode('utf-8'))

解説:
textractは多様なファイル形式（PDF, DOCX, PPTX, CSVなど）をテキストに変換するライブラリです。
変換後のテキストをrag.insert()で挿入可能です。

カスタムナレッジグラフの挿入

独自のKG(ナレッジグラフ)をJSON形式で挿入できます。
ここではNeo4Jなどの外部ストレージは使わず、デフォルト実装（NetworkX）を利用します。

カスタムKG挿入例（解説のみ）

# custom_kg = {
#     "entities": [
#         {
#             "entity_name": "CompanyA",
#             "entity_type": "Organization",
#             "description": "A major technology company",
#             "source_id": "Source1"
#         },
#         {
#             "entity_name": "ProductX",
#             "entity_type": "Product",
#             "description": "A popular product developed by CompanyA",
#             "source_id": "Source1"
#         }
#     ],
#     "relationships": [
#         {
#             "src_id": "CompanyA",
#             "tgt_id": "ProductX",
#             "description": "CompanyA develops ProductX",
#             "keywords": "develop, produce",
#             "weight": 1.0,
#             "source_id": "Source1"
#         }
#     ],
#     "chunks": [
#         {
#             "content": "ProductX, developed by CompanyA, has revolutionized the market.",
#             "source_id": "Source1",
#         },
#     ],
# }
#
# rag.insert_custom_kg(custom_kg)

解説:
insert_custom_kgを使えば、事前に定義したエンティティ・関係・チャンクをLightRAG内部のKGに統合できます。

不要なエンティティの削除

エンティティ削除（解説のみ）

# rag.delete_by_entity("Project Gutenberg")

解説:
delete_by_entityで特定エンティティ名を持つノードをKGから削除できます。

グラフ可視化（任意）

ここでは可視化は別途としますが、networkxやpyvisを用いてHTMLで可視化も可能です。

# ※本コードは実行例としてはコメントアウト
from pyvis.network import Network
import networkx as nx

G = nx.read_graphml('./dickens/graph_chunk_entity_relation.graphml')
net = Network(notebook=True)
net.from_nx(G)
net.show('knowledge_graph.html')

解説:
pyvisを用いると、生成されたナレッジグラフをHTML形式で可視化できますが、Neo4Jによる大規模可視化は今回は省略します。