DatabricksでLlamaIndexを動かしてみる

昨年このようなものがリリースされていたとは、自分はLangChainにかまけていました。

https://www.llamaindex.ai/

マニュアルはこちら。マニュアルによると、プライベートあるいはドメイン固有のデータの取り込み、構造化、アクセスのための[LLM](https://ja.wikipedia.org/wiki/%E5%A4%A7%E8%A6%8F%E6%A8%A1%E8%A8%80%E8%AA%9E%E3%83%A2%E3%83%87%E3%83%AB)ベースアプリケーションのためのデータフレームワークとのことです。

https://docs.llamaindex.ai/en/stable/index.html

こちらの記事でも説明されれています。

https://dev.classmethod.jp/articles/llamaindex-tutorial-001-overview-v0-7-9/

この記事ではまずは[チュートリアル](https://docs.llamaindex.ai/en/stable/getting_started/starter_example.html)をウォークスルーします。

最新のMLランタイムのクラスターを使います。
![Screenshot 2024-01-16 at 19.08.00.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/f5d05aed-7004-5f5d-9152-0d5b5af830ba.png)

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

```py
%pip install -U llama-index
dbutils.library.restartPython()
```

ドキュメントの準備。チュートリアルでは、[こちらのエッセイ](https://paulgraham.com/worked.html)を使っていますのでローカルにテキストファイルとして保存しておきます。あと、OpenAI APIを使うのでAPIキーを[シークレット](https://docs.databricks.com/ja/security/secrets/secrets.html)に保存しておきます。

そして、テキストの保存場所には[Unity Catalogのボリューム](https://docs.databricks.com/ja/connect/unity-catalog/volumes.html)を使います。ここでは、`takaakiyayoi_catalog.llama_index.data`というボリュームを作った後に以下を実行します。

```py
import os
os.environ["OPENAI_API_KEY"] = dbutils.secrets.get("demo-token-takaaki.yayoi", "openai_api_key")

from llama_index import VectorStoreIndex, SimpleDirectoryReader

TEXT_DIR = "/Volumes/takaakiyayoi_catalog/llama_index/data/corpus"
os.mkdir(TEXT_DIR)
```

すると`corpus`というフォルダができますので、先ほど保存したテキストファイルをアップロードします。
![Screenshot 2024-01-16 at 19.04.46.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/8cf2d35d-ac03-192c-89ea-6d3262ca296c.png)

以下を実行してベクトルインデックスを作成します。

```py
documents = SimpleDirectoryReader("/Volumes/takaakiyayoi_catalog/llama_index/data/corpus").load_data()
index = VectorStoreIndex.from_documents(documents)
```

クエリーを投げます。
```py
query_engine = index.as_query_engine()
response = query_engine.query("What did the author do growing up?")
print(response)
```
チュートリアルのものと似たような結果が返ってきます。

> The author mentioned that before college, the two main things they worked on outside of school were writing and programming. They wrote short stories and also tried writing programs on the IBM 1401 computer. They later got a microcomputer and started programming more extensively, writing simple games and a word processor.

ベクトルインデックスを永続化します。

```py
PERSIST_DIR = "/Volumes/takaakiyayoi_catalog/llama_index/data/indexes/tutorial"

# ドキュメントをロードしインデックスを作成
documents = SimpleDirectoryReader(TEXT_DIR).load_data()
index = VectorStoreIndex.from_documents(documents)
# 後で利用できるように保存
index.storage_context.persist(persist_dir=PERSIST_DIR)
```
保存されました。
![Screenshot 2024-01-16 at 19.09.42.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/a5e78abe-f968-e13d-948b-2fdaa5346b77.png)

```py
from llama_index import (
    VectorStoreIndex,
    SimpleDirectoryReader,
    StorageContext,
    load_index_from_storage,
)

storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
index = load_index_from_storage(storage_context)

# インデックスに対するクエリー
query_engine = index.as_query_engine()
response = query_engine.query("What did the author do growing up?")
print(response)
```

ベクトルインデックスを読み込んでクエリーを行うことができます。

```py
from llama_index import (
    VectorStoreIndex,
    SimpleDirectoryReader,
    StorageContext,
    load_index_from_storage,
)

storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
index = load_index_from_storage(storage_context)

# インデックスに対するクエリー
query_engine = index.as_query_engine()
response = query_engine.query("What did the author do growing up?")
print(response)
```

> The author mentioned that before college, the two main things they worked on outside of school were writing and programming. They wrote short stories and also tried writing programs on the IBM 1401 computer. They later got a microcomputer and started programming more extensively, writing simple games and a word processor.

OpenAIが動いているので、日本語でのクエリーもいけます。

```py
query_engine = index.as_query_engine()
response = query_engine.query("幼少期、著者は何をしましたか？")
print(response)
```

> 幼少期、著者は短編小説を書いたり、IBM 1401というコンピュータでプログラムを試したりしました。また、友人と一緒にヒースキットのキットでマイクロコンピュータを組み立てたりもしました。その後、父親にTRS-80というコンピュータを買ってもらい、プログラミングを本格的に始めました。

これは便利。RAG(Retrieval Augumented Generation)のアプリを簡単に動かせますね。

### はじめてのDatabricks

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

### Databricks無料トライアル

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