見出し画像

githubリポジトリをEmbeddingして質問に答えてもらう

ニケちゃん

先日、からあげ先生がこのようなツイートをされていました。任意のリポジトリから必要な情報がすぐ取得できたら嬉しいですよね。

面白そうだったので、実行した方法と試してみた結果を書きます。

ちなみに再現方法もからあげ先生が提示してくれたままの方法です。


やり方

上記のツイートのリンクにアクセスします。先に答えておくと、実態は「LlamaIndex」と「LlamaHub」が提供しているデータコネクタです。

必要なライブラリをインストールしておきます。

pip install llama_index llama_hub

リンク先の一番下のコードをコピペします。

import pickle
import os

from llama_index import download_loader, GPTVectorStoreIndex
download_loader("GithubRepositoryReader")

from llama_hub.github_repo import GithubClient, GithubRepositoryReader

docs = None
if os.path.exists("docs.pkl"):
    with open("docs.pkl", "rb") as f:
        docs = pickle.load(f)

if docs is None:
    github_client = GithubClient(os.getenv("GITHUB_TOKEN"))
    loader = GithubRepositoryReader(
        github_client,
        owner =                  "jerryjliu",
        repo =                   "llama_index",
        filter_directories =     (["gpt_index", "docs"], GithubRepositoryReader.FilterType.INCLUDE),
        filter_file_extensions = ([".py"], GithubRepositoryReader.FilterType.INCLUDE),
        verbose =                True,
        concurrent_requests =    10,
    )

    docs = loader.load_data(branch="main")

    with open("docs.pkl", "wb") as f:
        pickle.dump(docs, f)

index = GPTVectorStoreIndex.from_documents(docs)

query_engine = index.as_query_engine()
response = query_engine.query("Explain each LlamaIndex class?")
print(response)

オプション

下記を参考に設定してみてください。ちなみに、llama_indexのオーナー名が "jerryjliu" から "run-llama" に変更されているので、上記のコードのままだとエラーになります。

owner: リポジトリのオーナー名
repo: リポジトリ名
filter_directories: ルートから見てどのフォルダを対象にするか、第二引数は GithubRepositoryReader.FilterType.EXCLUDE か GithubRepositoryReader.FilterType.INCLUDE から選択します。EXCLUDEを選択した場合は、第一引数以外のファイルを対象にします。INCLUDEの場合は、第一引数のファイルのみを対象にします。全対象にしたい場合は、EXCLUDEを選択した上で、第一引数を[]にしてください。
filter_file_extensions: どの拡張子のファイルを対象にするかを決めます。第二引数はfilter_directoriesと同じです。
verbose: 処理中の詳細なログや情報が表示するかどうか
concurrent_requests: 同時に行うリクエストの最大数

最後に必要なトークン、キーを設定すれば完了です。

export OPENAI_API_KEY='...'
export GITHUB_TOKEN='...'

GITHUB_TOKENは下記の方法で取得できます。難しくないです。

準備が出来たら実行してみましょう。
質問は下記の部分で変更します。今回はデフォルトの質問に日本語で返すように追記しただけです。

response = query_engine.query("Explain each LlamaIndex class? in Japanese")

初めて実行する場合はリポジトリの読み込みで時間がかかります。

LlamaIndexの各クラスを説明します。

1. ObjectRetriever（オブジェクトリトリーバー）：オブジェクトを取得するための基本的な機能を提供するクラスです。

2. ObjectIndex（オブジェクトインデックス）：オブジェクトをインデックスするための基本的な機能を提供するクラスです。

3. SimpleObjectNodeMapping（シンプルオブジェクトノードマッピング）：オブジェクトとノードの間のシンプルなマッピングを提供するクラスです。

4. SimpleToolNodeMapping（シンプルツールノードマッピング）：ツールとノードの間のシンプルなマッピングを提供するクラスです。

5. SimpleQueryToolNodeMapping（シンプルクエリツールノードマッピング）：クエリツールとノードの間のシンプルなマッピングを提供するクラスです。

6. SQLTableNodeMapping（SQLテーブルノードマッピング）：SQLテーブルとノードの間のマッピングを提供するクラスです。

7. SQLTableSchema（SQLテーブルスキーマ）：SQLテーブルのスキーマを定義するクラスです。

8. KeywordTableIndex（キーワードテーブルインデックス）：キーワードテーブルをインデックスするための基本的な機能を提供するクラスです。

9. SimpleKeywordTableIndex（シンプルキーワードテーブルインデックス）：シンプルなキーワードテーブルをインデックスするためのクラスです。

10. RAKEKeywordTableIndex（RAKEキーワードテーブルインデックス）：RAKEアルゴリズムを使用してキーワードテーブルをインデックスするためのクラスです。

11. SummaryIndex（サマリーインデックス）：サマリーをインデックスするためのクラスです。

12. TreeIndex（ツリーインデックス）：ツリーをインデックスするためのクラスです。

13. VectaraIndex（Vectaraインデックス）：Vectaraを使用してインデックスを管理するためのクラスです。

14. DocumentSummaryIndex（ドキュメントサマリーインデックス）：ドキュメントのサマリーをインデックスするためのクラスです。

15. GPTKeywordTableIndex（GPTキーワードテーブルインデックス）：GPTを使用してキーワードテーブルをインデックスするためのクラスです（レガシー）。

16. GPTSimpleKeywordTableIndex（GPTシンプルキーワードテーブルインデックス）：GPTを使用してシンプルなキーワードテーブルをインデックスするためのクラスです（レガシー）。

17. GPTRAKEKeywordTableIndex（GPT-RAKEキーワードテーブルインデックス）：GPT-RAKEを使用してキーワードテーブルをインデックスするためのクラスです（レガシー）。

18. GPTListIndex（GPTリストインデックス）：GPTを使用してリストをインデックスするためのクラスです（レガシー）。

19. GPTTreeIndex（GPTツリーインデックス）：GPTを使用してツリーをインデックスするためのクラスです（レガシー）。

20. ListIndex（リストインデックス）：リストをインデックスするためのクラスです。

いけました！！

改造

このままだと使いづらいので少し改造します。

import os
import pickle

from llama_index import download_loader, GPTVectorStoreIndex
download_loader("GithubRepositoryReader")

from llama_hub.github_repo import GithubClient, GithubRepositoryReader

repository = input("Enter the repository name: ")
if repository == "":
    repository = "run-llama/llama_index"

query = ""
while query == "":
    query = input("Enter the query: ")
    if query == "":
        print("Please enter a query.")

owner, repo = repository.split("/")
file_name = f"{owner}_{repo}.pkl"

# docsフォルダが存在しなければ作成
if not os.path.exists('docs'):
    os.makedirs('docs')

# docsの読み込みまたは作成
docs_file_name = f"docs/{file_name}"
if os.path.exists(docs_file_name):
    with open(docs_file_name, "rb") as f:
        docs = pickle.load(f)
else:
    github_client = GithubClient(os.getenv("GITHUB_TOKEN"))
    loader = GithubRepositoryReader(
        github_client,
        owner=owner,
        repo=repo,
        filter_directories=(["llama_index", "docs"], GithubRepositoryReader.FilterType.INCLUDE),
        filter_file_extensions=([".py"], GithubRepositoryReader.FilterType.INCLUDE),
        verbose=True,
        concurrent_requests=10,
    )

    docs = loader.load_data(branch="main")

    with open(docs_file_name, "wb") as f:
        pickle.dump(docs, f)

# LLM定義
llm = OpenAI(model="gpt-4")
service_context = ServiceContext.from_defaults(llm=llm)

# indexの読み込みまたは作成
if os.path.exists(f"./storage/{repo}/vector_store.json"):
    storage_context = StorageContext.from_defaults(persist_dir=f"./storage/{repo}")
    index = load_index_from_storage(storage_context, service_context=service_context)
else:
    if not os.path.exists(f"./storage/{repo}"):
        os.makedirs(f"./storage/{repo}")

    set_global_service_context(service_context)
    index = VectorStoreIndex.from_documents(docs)
    index.storage_context.persist(persist_dir=f"./storage/{repo}")

query_engine = index.as_query_engine()
response = query_engine.aquery(query)
print("\n=== query result ===")
print(response)
# nowから経過秒数を計算
print("\n=== time ===")
print((datetime.datetime.now() - now).seconds)

docsファイルの名称を『docs/オーナー名+リポジトリ名」で保存するようにしました。

indexを『storage/リポジトリ名/**.json』で保存するようにしました。

LLMを定義できるようにしました。設定方法はLlamaIndexの仕様に合わせてください。

ターミナルでリポジトリ名とクエリ（質問）を入力できるようにしました。フォルダや拡張子指定はデフォルトのままなので適宜変更してください。

下記はindex保存後の検索結果です。初回は2分ほどかかりますが、以降は15秒で内容を損なわず結果を返却できています（GPT-3.5）。

-> % python github_repo.py
Enter the repository name: run-llama/llama_index
Enter the query: Explain each LlamaIndex class?
=== query result ===
The LlamaIndex package contains several classes. Here is an explanation of each class:

1. ObjectRetriever: This class is responsible for retrieving objects from the LlamaIndex.

2. ObjectIndex: This class represents an index of objects in the LlamaIndex.

3. SimpleObjectNodeMapping: This class provides a simple mapping between objects and nodes in the LlamaIndex.

4. SimpleToolNodeMapping: This class provides a simple mapping between tools and nodes in the LlamaIndex.

5. SimpleQueryToolNodeMapping: This class provides a simple mapping between query tools and nodes in the LlamaIndex.

6. SQLTableNodeMapping: This class provides a mapping between SQL tables and nodes in the LlamaIndex.

7. SQLTableSchema: This class represents the schema of an SQL table in the LlamaIndex.

8. DocumentSummaryIndex: This class represents an index of document summaries in the LlamaIndex.

9. KeywordTableIndex: This class represents an index of keyword tables in the LlamaIndex.

10. GPTKeywordTableIndex: This class represents a GPT-based keyword table index in the LlamaIndex.

11. RAKEKeywordTableIndex: This class represents a RAKE-based keyword table index in the LlamaIndex.

12. SimpleKeywordTableIndex: This class represents a simple keyword table index in the LlamaIndex.

13. SummaryIndex: This class represents an index of summaries in the LlamaIndex.

14. TreeIndex: This class represents a tree index in the LlamaIndex.

15. VectaraIndex: This class represents a Vectara index in the LlamaIndex.

16. GPTSimpleKeywordTableIndex: This class represents a GPT-based simple keyword table index in the LlamaIndex.

17. GPTRAKEKeywordTableIndex: This class represents a GPT-based RAKE keyword table index in the LlamaIndex.

18. GPTListIndex: This class represents a GPT-based list index in the LlamaIndex.

19. GPTTreeIndex: This class represents a GPT-based tree index in the LlamaIndex.

20. ListIndex: This class represents a list index in the LlamaIndex.

These classes provide various functionalities for managing and accessing different types of data in the LlamaIndex package.

=== time ===
120
-> % python github_repo.py
Enter the repository name: run-llama/llama_index
Enter the query: Explain each LlamaIndex class?
=== query result ===
The LlamaIndex package contains several classes. Here is an explanation of each class:

1. ObjectRetriever: This class is responsible for retrieving objects from the LlamaIndex.

2. ObjectIndex: This class represents an index of objects in the LlamaIndex.

3. SimpleObjectNodeMapping: This class provides a simple mapping between objects and nodes in the LlamaIndex.

4. SimpleToolNodeMapping: This class provides a simple mapping between tools and nodes in the LlamaIndex.

5. SimpleQueryToolNodeMapping: This class provides a simple mapping between query tools and nodes in the LlamaIndex.

6. SQLTableNodeMapping: This class provides a mapping between SQL tables and nodes in the LlamaIndex.

7. SQLTableSchema: This class represents the schema of an SQL table in the LlamaIndex.

8. DocumentSummaryIndex: This class represents an index of document summaries in the LlamaIndex.

9. KeywordTableIndex: This class represents an index of keyword tables in the LlamaIndex.

10. RAKEKeywordTableIndex: This class represents an index of keyword tables generated using the RAKE algorithm in the LlamaIndex.

11. SimpleKeywordTableIndex: This class represents a simple keyword table index in the LlamaIndex.

12. SummaryIndex: This class represents an index of summaries in the LlamaIndex.

13. TreeIndex: This class represents a tree index in the LlamaIndex.

14. VectaraIndex: This class represents a Vectara index in the LlamaIndex.

15. GPTKeywordTableIndex: This class represents a keyword table index generated using GPT in the LlamaIndex.

16. GPTSimpleKeywordTableIndex: This class represents a simple keyword table index generated using GPT in the LlamaIndex.

17. GPTRAKEKeywordTableIndex: This class represents a keyword table index generated using GPT and RAKE in the LlamaIndex.

18. GPTListIndex: This class represents a list index generated using GPT in the LlamaIndex.

19. GPTTreeIndex: This class represents a tree index generated using GPT in the LlamaIndex.

20. ListIndex: This class represents a list index in the LlamaIndex.

These classes provide various functionalities for managing and accessing different types of data in the LlamaIndex package.

=== time ===
15

streaming対応

最後の方をこのように変更してください。

query_engine = index.as_query_engine(streaming=True)
response = query_engine.query(query)
print("\n=== query result ===")
for text in response.response_gen:
    print(text)
-> % python test.py
Enter the repository name: run-llama/llama_index
Enter the query: このリポジトリの概要を教えてください

=== query result ===

この
リ
ポ
ジ
ト
リ
は
、
ベ
ク
ト
ル
スト
ア
と
評
価
ベ
ン
チ
マ
ー
ク
を
含
む
一
連
の
Python
ス
ク
リ
プ
ト
を
提
供
して
います
。
ベ
ク
ト
ル
スト
ア
に
は
、
"
Vector
Store
Query
"
、
"
Vector
Store
Query
Result
"
、
"
Metadata
Filters
"
な
ど
の
ク
ラ
ス
が
含
ま
れ
て
います
。
評
価
ベ
ン
チ
マ
ー
ク
に
は
、
"
hot
pot
qa
.py
"
と
いう
名
前
の
ファ
イ
ル
が
あり
、
Hot
pot
Q
AE
valuator
と
いう
ク
ラ
ス
を
含
ん
で
います
。
この
ク
ラ
ス
は
、
Hot
pot
QA
デ
ータ
セ
ット
を
ダ
ウ
ン
ロ
ード
し
、
それ
を
評
価
する
た
め
の
メ
ソ
ッ
ド
を
提
供
します
。

おまけ

-> % python github_repo.py
Enter the repository name: 
Enter the query: GithubRepositoryReaderクラスの各引数を説明してください。
=== query result ===
GithubRepositoryReaderクラスの各引数を説明します。

- owner (str): リポジトリの所有者を指定します。
- repo (str): リポジトリの名前を指定します。
- use_parser (bool): ファイルからテキストを抽出するためにパーサーを使用するかどうかを指定します。
- verbose (bool): 冗長なメッセージを表示するかどうかを指定します。
- github_token (str): Githubのトークンを指定します。指定しない場合、GITHUB_TOKEN環境変数から読み取られます。
- concurrent_requests (int): Github APIへの同時リクエスト数を指定します。
- ignore_file_extensions (List[str]): 無視するファイルの拡張子のリストを指定します。
- ignore_directories (List[str]): 無視するディレクトリのリストを指定します。

=== time ===
9

この記事が気に入ったらサポートをしてみませんか？