[LlamaIndex]AWS Lambdaで作成済みのインデックスをクエリしてみた

はじめに

こんにちは、筧( @TakaakiKakei )です。

LlamaIndexはOpenAIのLLMに独自のデータを読み込ませる仕組みです。 DevelopersIOでは既に色んなAWS環境で動かす方法が紹介されています。

本ブログではAWS Lambdaで作成済みのインデックスをクエリする例をご紹介します。 リアルタイム性をそこまで求められない場合に、役に立つかもしれません。

コード

Serverless Framework を利用してデプロイします。 デプロイ方法はREADMEを参照ください。

工夫点

レスポンス改善

handler内でインデックスのロード処理を書くと、Lambdaが実行する度にロードしてしまうため、レスポンスが遅くなるという問題がありました。 ロード時間の目安は以下の記事内の結果を参照ください。

そこで以下の2点で、レスポンス改善をしました。

• handler外にロード処理を書くことで、コールドスタートでインデックスをロード
• Provisioned Concurrencyの設定を追加することで、コールドスタートの頻度をへらす

src/query.py

"""
- 回答を早く返すために、コールドスタンバイで、インデックスをロードしておく
    - デプロイ前に、環境変数 INDEX_BUCKER_NAME の S3バケットに index.json を配置する必要あり
    - インデックスの更新の度に再デプロイが必要
"""
openai.api_key = os.environ["OPENAI_API_KEY"]

# Lambdaの/tmp配下にインデックスを格納
tmp_index_file_name = (
    "/tmp/index-"
    + datetime.now().strftime("%Y-%m-%d-%H-%M-%S-")
    + str(random.randint(0, 999999))
    + ".json"
)
s3 = boto3.client("s3")
s3.download_file(os.environ["INDEX_BUCKER_NAME"], "index.json", tmp_index_file_name)
# index.jsonをロード
llm_predictor = LLMPredictor(
    llm=OpenAI(temperature=0.7, model_name="gpt-3.5-turbo", max_tokens=512)
)
loaded_index = GPTSimpleVectorIndex.load_from_disk(
    save_path=tmp_index_file_name,
    llm_predictor=llm_predictor,
)
# Lambdaの/tmp配下に格納したインデックスを削除
os.remove(tmp_index_file_name)

src/serverless.yml

functions:
  query:
    handler: src/query.handler
    ephemeralStorageSize: 10240
    provisionedConcurrency: 1

インデックスの読み込み

インデックスをS3からAWS Lambdaのエフェメラルストレージにダウンロードすることで、インデックスのサイズが大きくなっても対応できるようにしました。

src/serverless.yml

functions:
  query:
    handler: src/query.handler
    ephemeralStorageSize: 10240

外部モジュールの読み込み

llama-indexの読み込み方法を工夫しました。 はじめはpoetryを利用していましたが、AWS Lambda実行時に、Runtime.ImportModuleErrorが発生してしまいました。 具体的には、Importing the numpy C-extensions failed.という内容です。 根本解決はできていませんが、pipenvを利用することで回避できました。

pipenvの利用に切り替えたので、外部モジュールの読み込みには、dockerizePip: trueを利用しました。 またllama-indexの容量が大きかったので、zip: true, slim: true を利用しました。

src/serverless.yml

custom:
  pythonRequirements:
    dockerizePip: true
    zip: true
    slim: true

動かしてみる

devステージでデプロイしていると、llamaindex-lambda-query-dev というステートマシンができています。 こちらを使って実行してみます。

無事、LlamaIndexを使ってクエリできました。

おわりに

最後まで読んでいただきありがとうございます。

現状、回答の精度に課題がありそうですが、今後調査して改善できればと思います。

それではまた！