Llama-indexのIndexをS3に保存する

2023年07月20日 1分で読めます

新技術の試用 • DIYコーディング

Llama-indexのIndexをS3に保存したい。ドキュメントとしてはここに記載がある。

https://gpt-index.readthedocs.io/en/latest/core_modules/data_modules/storage/save_load.html#using-a-remote-backend

s3fs を使っている。

`s3fs` とは

S3をファイルシステムのように扱うライブラリ。fsspecをベースにしている。

https://pypi.org/project/s3fs/

使ってみよう。

前提として以下を設定する。

AWSにIAMユーザーを作成する。ACCESS KEY ID, ACCESS KEY SECRETを取得する
AWSにS3バケットを作成する。バケット名を確認しておく。

使ってみるとこうなる

>>> import s3fs
>>> aws_key = "..." # AWSで利用するACCESS KEY IDを設定する。
>>> aws_secret = "..." # AWSで利用するACCESS KEY IDのSECRETを設定する
>>> bucket_name = "..." # 作成したバケット名を設定する
>>>
>>> s3 = s3fs.S3FileSystem(key=aws_key, secret=aws_secret)
>>> s3.ls(bucket_name)
[]
>>> with s3.open(f'{bucket_name}/new-file.txt', 'w') as f:
...     f.write("Hello s3fs!!")
...
12
>>> s3.ls(bucket_name)
['.../new-file.txt'] # バケット名の部分はマスキングしてる
>>> with s3.open(f'{bucket_name}/new-file.txt', 'r') as f:
...     data = f.read()
...
>>> print(data)
Hello s3fs!!

llama-indexのIndexをS3に保存する

上記のs3fsを利用してS3にインデックスを保存する。今回ははてなブログから記事を読むことにした。保存して読み込むコードはこんな感じ

import s3fs
from decouple import config
from langchain.chat_models import ChatOpenAI
from llama_hub.hatena_blog.base import HatenaBlogReader
from llama_index import (
    GPTVectorStoreIndex,
    LLMPredictor,
    PromptHelper,
    ServiceContext,
    StorageContext,
    load_index_from_storage,
)


# はてなブログから記事を読み込む
reader = HatenaBlogReader(
    root_endpoint=config("ATOM_PUB_ROOT_ENDPOINT"),
    api_key=config("ATOM_PUB_API_KEY"),
    username=config("HATENA_BLOG_USERNAME"),
)
documents = reader.load_data()

# LLMのコンテキストを設定
llm_predictor = LLMPredictor(llm=ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo"))
prompt_helper = PromptHelper(context_window=4096, num_output=256)
service_context = ServiceContext.from_defaults(
    llm_predictor=llm_predictor, prompt_helper=prompt_helper
)

# インデックス作成
index = GPTVectorStoreIndex.from_documents(
    documents, service_context=service_context
)

# インデックス保存
fs = s3fs.S3FileSystem(key=config("AWS_ACCESS_KEY_ID"), secret=config("AWS_SECRET_ACCESS_KEY"))
index.storage_context.persist(persist_dir=f{bucket_name}/index/hatena", fs=fs)

# インデックスを読み込む
storage_context = StorageContext.from_defaults(persist_dir=f{bucket_name}/index/hatena", fs=fs)
index = load_index_from_storage(storage_context=storage_context, service_context=service_context)

# クエリを実行する
query_engine = self.index.as_query_engine()
response = query_engine.query(message)

print(response.response)

s3fs とは

llama-indexのIndexをS3に保存する

`s3fs` とは