CCCMKホールディングス TECH LABの Tech Blog

TECH LABのエンジニアが技術情報を発信しています

ブログタイトル

トップ > LLM > MLflow 2.3で搭載されたLLMの管理機能を試してみました！

MLflow 2.3で搭載されたLLMの管理機能を試してみました！

LLM MLflow

こんにちは、CCCMKホールディングス TECH LABの三浦です。

5月はずっと好きなアーティストのアルバムがリリースされたり、楽しみにしていたゲームが発売されたり、うれしいことがたくさんありました。心に栄養を取ってる感じがしました。

先日Databricksのウェビナー"Hello Dolly オープンな大規模言語モデル(LLM) を自身でトレーニグ"を拝聴しました。とても刺激が多いウェビナーで、自分でも試してみたい、という内容がたくさん見つかりました。

その中でもすぐに試そう、と思ったのがMLflow 2.3によるLLMの管理運用機能です。MLflowのModel Registryでは機械学習モデルを管理し、それらを共通化された手順で様々な環境にデプロイすることが出来ます。PyTorchやScikit-Learn等のフレームワークに加え、MLflow2.3からはLLMに関連するHugging Face Transformers、OpenAI関数、LangChainも対応が開始されました。

www.databricks.com

Azure OpenAIをLangChainを通して使うと色々なことが出来ることは分かってきたものの、サービスとして活用するためにはプログラムを管理し、デプロイする手順を考えないといけません。MLflow 2.3でのLangChain対応はまだ実験的な段階ですが、今のうちに触れておきたい、と思い、試してみることにしました。

今回はAzure OpenAIとLangChainを使い、入力されたテキストをハッシュタグ化する処理を作り、それをMLflowのModel Registryに登録、別の環境からダウンロードして利用する、という流れを試してみました。

Model Registryに登録するまで

この部分はAzure DatabricksのNotebook上で試しています。使用したDatabricks Runtimeのバージョンは"13.0 ML"です。コードは先ほどのDatabricksのMLflow 2.3のブログを参考にさせて頂きました。

セットアップ

最初に必要になるライブラリをインストールします。

%pip install mlflow==2.3 openai==0.27.6 langchain==0.0.142

ちょうどこの実験をしていた時にopenaiの0.27.7が公開されたのですが、そちらを使うと私の環境ではAzure OpenAIのモデルにアクセスする時に"Resource Not Found"のエラーが出てしまい、これまで動いていたコードが動かなくなりました。原因を調べないと・・・と思いつつ、今回は取り急ぎ稼働実績のある0.27.6を使うようにして対処しました。

このコマンドを実行すると、Notebook上に"Note: you may need to restart the kernel using dbutils.library.restartPython() to use updated packages."というメッセージが表示されました。メッセージの指示通り、次に以下のコマンドを実行しました。これを実行しないと使用されるMLflowのバージョンが2.3に切り替わらないようです。

 dbutils.library.restartPython()

LLMChainを作る

MLflow 2.3で管理対象になるのはLangChainにおいてはLLMChainクラスで実装したモデルだけのようです。LLMChainを使ってメインの処理を書いていきます。

import mlflow
from langchain.llms import OpenAI
from langchain import LLMChain
from langchain.prompts import PromptTemplate

prompt_template = """
ユーザーが入力した全てのテキストに対し、ハッシュタグを抽出してください。

入力テキスト:{input_text}
ハッシュタグ:
"""

prompt = PromptTemplate(template=prompt_template, input_variables=['input_text'])

#AzureOpenAIではなくOpenAIを使用する
llm = OpenAI(
  engine="text-davinci-003", #deployment_nameだとだめ
  temperature=0
)

llm_chain = LLMChain(
  llm=llm,
  prompt=prompt
)

この部分には注意点があります。LLMChainオブジェクトを作るとき、Azure OpenAIのモデルを使う場合は引数llmにlangchain.llms.AzureOpenAIを指定するのですが、今のところMLflowではHuggingFaceHubもしくはOpenAIでllmを指定する必要があるようです。Model Registryに登録する時には警告が表示されるものの登録はされます。しかしModel Registryからロードする時にエラーが発生し、モデルを使うことが出来なくなりました。それからAzure OpenAIのdeploymentを指定する時はdeployment_nameではなく、engineで指定する必要があります。

Model Registryに登録する

先に作ったLLMChainのオブジェクトllm_chainの動作を確かめておきます。

llm_chain.run({'input_text':'今日はLLMについて話します。MLflow2.3がリリースされ、その中でLLMのサポートが始まりました。'})

出力結果

'#LLM #MLflow #MLflow2.3'

良さそうなので、これをModel Registryに登録します。

mlflow.langchain.log_model(
  lc_model=llm_chain,
  artifact_path='model',
  registered_model_name='extract-hashtag-with-langchain'
)

処理が完了したら、Model Registryに登録されていることが確認出来ます！

Model Registryからロードする

Azure DatabricksのNotebookから

先ほど登録したモデルをダウンロードして利用する手順は以下の様になります。モデルを利用する環境も上と同じライブラリをインストールし、Azure OpenAIへの接続情報を環境変数にセットしておく必要があります。

import mlflow

infer_model = mlflow.pyfunc.load_model(
  'models:/extract-hashtag-with-langchain/6',
)

infer_model.predict([{'input_text':'今日はLLMについて話します。MLflow2.3がリリースされ、その中でLLMのサポートが始まりました。'}])

出力結果

['#LLM #MLflow #MLflow2.3']

簡単な手続きでモデルをダウンロードして利用することが出来ました！

別の環境で動くDocker Containerから

次はもう少し実稼働をイメージした利用方法を試してみます。最近私はDockerの勉強中で、Docker ContainerからModel Registryのモデルをダウンロードして使う方法を調べてみました。

Containerで動くアプリケーションは、PythonのFastAPIを使って以下の様に書いてみました。("main.py")

from pydantic import BaseModel
import mlflow
from fastapi import FastAPI

class Text(BaseModel):
    """
    /tags/で受け付けるRequest Body
    """
    input_text: str

app = FastAPI()
infer_model = mlflow.pyfunc.load_model(
  'models:/extract-hashtag-with-langchain/0') #Databricks Model Registryから

@app.get("/")
async def root():
    return {"message": "テキストからタグを抽出します。"}

@app.post("/tags/")
async def create_tag(text: Text):
    print(text)
    tags = infer_model.predict(
        [{"input_text":text.input_text}])
    return tags

"/tags/"というエンドポイントでメッセージを{"input_text":"テキスト"}のように受け取ると、テキストからハッシュタグを抽出し、返してくれるアプリケーションです。

requirements.txtの内容は以下の通りです。

mlflow==2.3
openai==0.27.6
langchain==0.0.142
fastapi[all]

最後にDockerfileです。

FROM python:3.9-slim

WORKDIR api
COPY requirements.txt .
COPY mycert.pem mycert.crt
RUN pip install \
    --cert mycert.crt\
    -r requirements.txt
COPY main.py .
EXPOSE 8000
CMD [ "uvicorn","main:app","--host", "0.0.0.0", "--port", "8000" ]

(私のローカル環境ではpipを実行する時にSSLエラーが発生するため、SSL証明書を指定しています。)

imageをbuildします。

docker build -t using-llm-fastapi:0 .

Azure OpenAIやMLflowに接続するための接続情報はContainer起動時に指定します。それらは一旦ファイルにまとめるようにしました。

OPENAI_API_KEY=....
OPENAI_API_BASE=https://.../
OPENAI_API_VERSION=...
OPENAI_API_TYPE=azure
MLFLOW_TRACKING_URI=databricks
DATABRICKS_HOST=https://.../
DATABRICKS_TOKEN=...
REQUESTS_CA_BUNDLE=mycert.crt

(こちらでもSSLエラーを回避するためSSL証明書を指定しています。) 最後にContainerを起動します。

 docker run --rm -it -p 8000:8000/tcp\
 --env-file env_file.txt\
 using-llm-fastapi:0

起動を確認したら、別のターミナルからリクエストを送信してみます。

curl -X POST -H "Content-Type:application/json" \
  -d '{"input_text":"PythonのFastAPIを使ってLLMのモデルをうごかしてみました。"}' \
  http://localhost:8000/tags/

すると以下の様に応答が得られました。ちゃんと動いてますね・・・！

["#Python #FastAPI #LLM #モデル"]

まとめ

ということで、今回はMLflow2.3で搭載されたLLMモデルの管理機能について試してみたことをまとめてみました。MLflowを使うことでLLMを使った開発からデプロイまでがだいぶ分かりやすくなり、シェアもしやすくなると思いました。Azure OpenAIやLangChainへの対応は今後進んでいくと思います。もっと効率的にLLMを扱えるよう、今からMLflowを活用した管理運用方法を考えていければ・・・と思いました。