Ollamaを使って推論してみる on Databricks

前から使ってみたいと思いつつ、ようやく試せました。

# 導入

ローカルLLMの推論パッケージのひとつに[Ollama](https://github.com/ollama/ollama)があります。

https://ollama.com/

かなり簡単にモデルのロードや推論が実行できるパッケージであり、またGGUFで量子化されたモデルの読み込みもサポートしています。

LangChainの[LangChain State of AI 2023](https://blog.langchain.dev/langchain-state-of-ai-2023/)でもMost Used OSS Model ProvidersやMost Used LLM Providersの上位に位置しており、非常によく利用されているパッケージのひとつです。

Databricks上でGGUFフォーマットを手軽に使えるパッケージが欲しく、前から注目はしていたのですが、試す時間ができたので触ってみました。

検証はDatabricks on AWS上で実施しました。
DBRは14.3ML、クラスタタイプはg4dn.xlarge(GPUクラスタ)を利用しています。

# Step0. サーバ用のバイナリファイルをダウンロード

Ollamaサーバ用のバイナリファイルをダウンロード・インストールします。

標準ではcurlからshellを実行する方法が示されていますが、自分の環境下では問題があったため、バイナリファイルをダウンロードして使うことにします。

[こちら](https://ollama.com/download/ollama-linux-amd64)にバイナリファイルがあるので、ここからファイルをダウンロードし、Unity Catalog Volumes上にアップロード。

場所は適当に用意した```training.llm.tmp```ボリュームにしました。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/83c1f96a-7672-fa3c-a6d9-32922734febc.png)


# Step1. パッケージのインストールなど

ノートブックを作成し、Python用のパッケージをインストールします。
これはPython用のクライアントライブラリとなります。

```python
%pip install ollama

dbutils.library.restartPython()
```

次に先ほどUnity Catalog Volumesにアップロードしたバイナリファイルをクラスタ側にコピー。

```python
dbutils.fs.cp("/Volumes/training/llm/tmp/ollama-linux-amd64", "file:/usr/bin/ollama")
```

実行権限を付与します。

```
!chmod +x /usr/bin/ollama
```

# Step2. Ollamaサーバを起動

ノートブック画面下部からターミナルを開き、以下のコマンドを実行してOllamaサーバを起動します。

```
/usr/bin/ollama serve
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/c1cf6c1f-ccdc-4949-2d78-4e9d4fd37021.png)

# Step3. モデルファイルを取得(Pull)

ノートブック上で以下のセルを実行し、モデルファイルを取得します。
今回はGoogle Gemma-7b-itを利用しました。

https://ollama.com/library/gemma

モデルファイルはOllama用のリポジトリから取得している？ようです。
公式リポジトリでは(おそらく)GGUFで4bit量子化されたものがデフォルトで提供されています。（それ以外のbitで量子化されたものも提供されています）

```python
import ollama
ollama.pull('gemma:7b-instruct')
```

# Step4. 推論

準備が整ったので推論してみます。
まずは公式サンプルで単純なChat生成。

```python
import ollama
response = ollama.chat(model='gemma:7b-instruct', messages=[
  {
    'role': 'user',
    'content': 'Why is the sky blue?',
  },
])
print(response['message']['content'])
```

```:出力
The sky appears blue due to a phenomenon called **Rayleigh Scattering**.

**Rayleigh Scattering:**

* Sunlight consists of various colors of the spectrum, including red, green, blue, and violet.
* When sunlight enters the Earth's atmosphere, particles of air (dust, smoke, water droplets) scatter the different colors of light in different directions.
* Scattered light travels in all directions, but the scattered light in the direction of the observer's eyes is more prominent.

**Blue Sky:**

* The scattered light in the direction of the observer's eyes is mostly the blue light.
* This scattered light reaches our eyes from all directions in the sky, making it appear blue.
* The intensity of the blue color is highest at sunrise and sunset, when the sun's rays travel through a longer path in the atmosphere and scatter more effectively.

**Other Factors:**

* The intensity of the blue color is also influenced by the time of day, season, and atmospheric conditions.
* Clouds and pollution can reduce the amount of scattered blue light, making the sky appear white or gray.
* At night, the sky appears black because there is less scattered light from the atmosphere.

**Conclusion:**

The sky appears blue due to Rayleigh Scattering, where particles in the atmosphere scatter different colors of light in different directions. The scattered light in the direction of the observer's eyes is predominantly blue, making the sky appear blue.
```

出力できましたね。

<br>

ストリーミング出力もしてみます。

```python
import ollama

stream = ollama.chat(
    model="gemma:7b-instruct",
    messages=[{"role": "user", "content": "Databricksとは何?"}],
    stream=True,
    options={
        "temperature": 0,
        "num_predict": 512,
    },
)

for chunk in stream:
    print(chunk["message"]["content"], end="", flush=True)
```

```:出力
Databricksは、データ分析のプラットフォームを提供する企業です。オープンソースに基づいて、データ分析を簡単にするために設計されています。

**主な特徴:**

* **データ統合:** データを複数のソースから統合し、分析するための統合されたデータセットを提供します。
* **データ分析:** データ分析を簡単にするために、 SQL、Python、R、Scalaなどの言語をサポートします。
* **データビジュアル:** データ分析の結果を視覚的に表示し、見やすさを向上させます。
* **スケーラビリティ:** 大量のデータを処理するためのスケーラブルな設計です。
* **コスト削減:** 開発コストやサーバーコストを削減するためのクラウドベースのプラットフォームです。

**主な使用ケース:**

* **ビッグデータ分析:** 大量のデータを分析するためのツール
* **データマイニング:** データ分析に基づいて意思決定するためのツール
* **データ visualization:** データ分析の結果を視覚的に表示するためのツール
* **機械学習:** データ分析に基づいて機械学習モデルを構築するためのツール

**主なメリット:**

* **データ分析の簡易化:** データ分析を簡単にするため、ビジネスのあらゆるユーザーが活用できます。
* **コスト削減:** 開発コストやサーバーコストを削減するためのクラウドベースのプラットフォームです。
* **スケーラビリティ:** 大量のデータを処理するためのスケーラブルな設計です。
* **データ統合:** データを複数のソースから統合し、分析するための統合されたデータセットを提供します。
```

```stream```パラメータをTrue指定することでストリーミングのチャンクを取得できるようになります。
また、```temperature```や```num_predict```(いわゆるmax_tokens)を```options```パラメータに指定することで制御できます。

# まとめ

OllamaをDatabricks上で試してみました。

サーバがGoで作成されているようで、軽快に動作します。
モデルの取得なども非常に簡単で、クライアントパッケージのインストールもシンプル、チャットテンプレートも標準提供というシンプルかつ容易に扱えるパッケージだと感じました。
今回試してませんが、LangChainとの親和性も高いようなので、プロトタイプ構築をサクっとやる際などに非常によさそうです。

また、これで主要な量子化フォーマットであるGGUF、AWQ、GPTQ、EXL2それぞれをDatabricks上で扱える状態になりました。最近はGGUFとEXL2がメジャーになりつつある気がしていますが、量子化はメモリ節約や推論速度向上において重要な技術だと思いますので、今後もキャッチアップしていきたいと思います。
今後は[AQLM](https://github.com/Vahe1994/AQLM)を追ってみたい。