LangServeを触ってみる on Databricks

以前からやろうと思っていたのですが、ようやく触れました。

# 導入

LangServeは```langchain```のChainを簡単にREST APIとしてデプロイすることができる仕組です。
FastAPI等との統合がされており、APIドキュメントの生成や、APIのPlayground機能なども提供しています。

https://www.langchain.com/langserve

Langchain TemplateもLangServeで試すような構造になっています。
また、Langchain側でホストするLangServe環境も今後公開される予定のようです。

実際どんな感じなんだろうと思っていたので、ローカルLLMを使ってDatabricks上で実行してみます。

:::note info

Databricksにはモデルサービング機能がありますので、本番利用においてはこちらを利用する方がよいと思います。
今回は開発・検証用途という位置づけです。
:::



# Step1. サーバの構築

LangServを使ってREST APIサーバを構築します。

## インストール

まずは必要なモジュールのインストール。
AWQフォーマットのモデルを使うので、autoawqもインストールします。

```python
%pip install -U transformers accelerate langchain
%pip install fastapi uvicorn nest_asyncio "autoawq==0.1.5" "langserve[server]"

dbutils.library.restartPython()
```

Databricksのノートブック上でFastAPIを使えるようにするために、```nest_asyncio```パッケージでイベントループのネストを有効にします。(手前ミソですが、このあたりの補足は[こちら](https://qiita.com/isanakamishiro2/items/c55123343e95fd2769a6#fastapi%E3%82%B5%E3%83%BC%E3%83%90%E3%81%AE%E8%B5%B7%E5%8B%95)の記事に記載）

```python
import nest_asyncio
import uvicorn

nest_asyncio.apply()
```

## モデルダウンロード

利用するモデルをダウンロードします。
今回は現在Huggingface上で人気のあるopenchat_3.5のAWQ量子化モデルを利用します。
Mistral 7Bベースの多言語モデルですが、日本語もそれなりに使えて使い勝手が良さそうな印象。

https://huggingface.co/TheBloke/openchat_3.5-AWQ

```python
def download_model(model_id:str):
    import os
    from huggingface_hub import snapshot_download

    UC_VOLUME = "/Volumes/training/llm/model_snapshots"

    local_dir = f"/tmp/{model_id}"
    uc_dir = f"/models--{model_id.replace('/', '--')}"

    snapshot_location = snapshot_download(
        repo_id=model_id,
        local_dir=local_dir,
        local_dir_use_symlinks=False,
    )

    dbutils.fs.cp(f"file:{local_dir}", f"{UC_VOLUME}{uc_dir}", recurse=True)


model_id = "TheBloke/openchat_3.5-AWQ"
download_model(model_id)
```

## モデルのロード

先ほどダウンロードしたモデルをロードし、langchainのChat Modelとしてラップします。
```ChatHuggingFaceModel```クラスは、[こちらの記事](https://qiita.com/isanakamishiro2/items/66411e1443cf78a2c1e8#step4-langchain-%E3%82%AB%E3%82%B9%E3%82%BF%E3%83%A0%E3%83%81%E3%83%A3%E3%83%83%E3%83%88%E3%83%A2%E3%83%87%E3%83%AB%E3%81%AE%E6%BA%96%E5%82%99)で作成した自作クラスを再利用しました。
もちろんローカルLLMを使わずに```langchain.llms.openai.OpenAI```などと入れ替えてもらって大丈夫です。

```python

from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers_chat import ChatHuggingFaceModel

def get_downloaded_model_path(model_id:str) -> str:
    UC_VOLUME = "/Volumes/training/llm/model_snapshots"
    uc_dir = f"/models--{model_id.replace('/', '--')}"

    return f"{UC_VOLUME}{uc_dir}"

model_path = get_downloaded_model_path(model_id)

generator = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_path)

chat_model = ChatHuggingFaceModel(
    generator=generator,
    tokenizer=tokenizer,
    human_message_template="GPT4 Correct User: {}
",
    ai_message_template="GPT4 Correct Assistant: {}",
    repetition_penalty=1.2,
    temperature=0.1,
    max_new_tokens=1024,
)
```

## プロンプトテンプレートの作成

単純に与えられた内容をそのまま実行するテンプレートにします。

```python
from langchain.prompts import ChatPromptTemplate
from langchain.prompts.chat import (
    AIMessagePromptTemplate,
    HumanMessagePromptTemplate,
)

template = """{question}"""
prompt = ChatPromptTemplate.from_messages(
    [
        HumanMessagePromptTemplate.from_template(template),
        AIMessagePromptTemplate.from_template(""),
    ]
)
```

## チェーンの作成&LangServeの起動

ようやく今日のメイン。

FastAPIのオブジェクトを作成。

```python
from fastapi import FastAPI, Response, Request
from fastapi.responses import HTMLResponse
from langserve import add_routes

app = FastAPI(
    title="LangChain Server",
    version="1.0",
    description="A simple api server using Langchain's Runnable interfaces",
)
```

ルートの登録。
```add_routes```を呼び出し、実行するChain(第２引数)を登録します。
今回はプロンプトテンプレートとLLM(Chat Model)だけのシンプルなChainです。

```python
add_routes(
    app,
    prompt | chat_model,
    path="/chain",
)
```

最後にAPIサーバを起動。

```python
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)
```

# Step2. クライアント準備/リクエスト実行

別のノートブックを作成し、LangServeサーバと通信するクライアントを用意します。

## モジュールインストール

必要なモジュールをインストール。

```python
%pip install -U langchain "langserve[client]"

dbutils.library.restartPython()
```

APIサーバのURLを取得する関数を作成。
今回はサーバ側と同一のクラスタ上で動かすことを想定して、単純に```http://127.0.0.1:8000```を返すようにしています。別の環境でクライアントを実行する場合は、Driver Proxy APIを利用してください。(コメントアウト部分を外すと該当のURLが取得できます)

```python
from dbruntime.databricks_repl_context import get_context

def api_url(port: int = 8000):
    # API を実行するための URL を取得
    # ブラウザでノートブックを開いているとき、ノートブック上で実行される各プロセスで使用できる URL を生成
    proxy_url = f"http://127.0.0.1:{port}/"
    # ctx = get_context()
    # proxy_url = f"https://{ctx.browserHostName}/driver-proxy-api/o/{ctx.workspaceId}/{ctx.clusterId}/{port}/"

    return proxy_url
```

## リクエストを実行

通常の```requests```呼び出しでもよいのですが、今回はPython用のSDKを使ってリクエストを実行します。
詳細は下記リンク先を確認ください。

https://github.com/langchain-ai/langserve?ref=blog.langchain.dev#client

```RemoteRunnable```を使って、リモート接続用のChainを作成。


```python
from langchain.schema import SystemMessage, HumanMessage
from langchain.prompts import ChatPromptTemplate
from langchain.schema.runnable import RunnableMap
from langserve import RemoteRunnable

# Access Tokenの設定。ローカルアドレスとの通信においては無くても動く。
api_token = dbutils.secrets.get("llm", "api_token")
headers = {"Authorization": f"Bearer {api_token}"}

# リモート接続するチェーンを作成
chain = RemoteRunnable(api_url()+"chain", headers=headers)
```

あとは通常のChainと同様に使えます。
単純なリクエストを実行してみます。

```python
question = "Databricksとは何かを箇条書きで5個以内で回答してください。"
result = chain.invoke({"question": question})

print(result.content)
```

```:出力
1. Databricksは、アプリケーションを簡単に構築、デプロイ、スケールするためのデータ分析プラットフォームです。
2. データ分析のための統合環境を提供し、PythonやRを使用してデータを分析できます。
3. データをストレージし、処理するための高性能なクラスターを提供します。
4. 機械学習、データフロー、データベース、およびビジュアライザーを含む、統合されたデータ分析ツールを提供します。
5. クラウド上で実行されるため、スケーラビリティと柔軟性が高く、データ分析のスピードやスケーラビリティが向上します。
```

```stream```/```astream```によるストリーミングでの出力も可能。

```python
question = "Databricksとは何かを箇条書きで5個以内で回答してください。"

# Supports astream
async for msg in chain.astream({"question": question}):
    print(msg.content, end="", flush=True)
```

# まとめ

langchainのChainを簡単にREST API化できるのはいいと思いました。
ただ、環境起因なのかレスポンスが返ってこないなど動作が少し不安定な感じ。
Databricks上でも安定的に使えると、開発用のRESTサーバとかに容易に使えそう。