MLFLOWにおけるLangChainモデルのログ記録APIの動作確認 #langchain

はじめに

LLMの実行管理が可能なツールについて、LangSmithは一般的によく知られていますが、他にも利用できるツールはないか調査しました。

以前に画像分類モデルの実験管理にMLFLOWを利用した経験があり、そのMLFLOWがLLM向けのAPIがリリースしたことを知り、気になったため、試しに動かしてみました。

検証環境

動作検証に用いた環境は以下となります。

WSL2（Ubuntu22.04 LTS）
CPU：Intel(R) Core(TM) i7-9700
RAM：16.0 GB

環境構築

Dockerで環境構築を実施します。

実行環境の構築

MLFLOWサーバとPythonの実行環境を立ち上げます。

Python実行環境

DockerでMLFLOWとjupyterlabが利用できるコンテナを起動

dockerfile

# 公式のPython 3イメージをベースに使用
FROM python:3

# パッケージリストを更新し、必要なパッケージをインストール
RUN apt-get update && \
    apt-get install -y \
    && rm -rf /var/lib/apt/lists/*

# pipをアップグレードし、必要なPythonパッケージをインストール
RUN pip install --upgrade pip && \
    pip install \
    jupyterlab \
    mlflow \
    langchain \
    openai \
    chromadb \
    numexpr \
    google-search-results

dockercompose

compose.yml

version: '3'

# JupyterLabを実行するコンテナ
services:
  jupyterlab:
    build: .  # Dockerfileを元にイメージをビルド
    image: jupyterlab  # コンテナのイメージ名
    container_name: jupyterlab-mlflow  # コンテナの名前
    working_dir: /work  # コンテナ内での作業ディレクトリ
    volumes:  # ホストとコンテナのディレクトリをマウント
      - ./work:/work
    ports:  # ホストとコンテナのポートをマッピング
      - "8888:8888"  # JupyterLab
      - "5000:5000"  # MLflow
    command: ["jupyter-lab", "--ip", "0.0.0.0", "--allow-root", "-b", "localhost"]  # コンテナ起動時のコマンド

MLFLOWサーバの起動

チュートリアルを参照して、Python実行環境のコンテナ内でMLFLOWサーバを起動する

起動後の画面

任意のブラウザからhttp://localhost:5000でMLFLOWを起動することができます。

検証

mlflowが持つ、LangChain モデルのログ記録とロードのためのAPIを検証する

MLflowのlangchain APIの検証

以下のページを参考にMLflowのlangchain APIを検証する

実装：MLflowでのLangChain のログ記録

インターネット上（Googleカスタム検索連携）から回答を得て、MLflowでモデルと推論結果を管理するユースケースを想定して実装する。

ex1.ipynd

import os

from langchain.agents import AgentType, initialize_agent, load_tools
from langchain.llms import OpenAI

import mlflow

mlflow.set_tracking_uri(uri="http://localhost:5000/")

# 注意: SerpAPI と OpenAI のAPIを使用するためのアカウントを持っていることを確認してください。
# 必要なAPIキーが設定されていることを確認する
assert "OPENAI_API_KEY" in os.environ, "環境変数 OPENAI_API_KEY が設定されていません。"
assert "SERPAPI_API_KEY" in os.environ, "環境変数 SERPAPI_API_KEY が設定されていません。"

# Enable mlflow langchain autologging（自動でログ取得）
mlflow.langchain.autolog(
    log_input_examples=True, # Trueの場合、推論データからの入力例が収集され、Langchainモデルのアーティファクトとともに推論時にログに記録されます。
    log_model_signatures=True, # Trueの場合、モデルの入力と出力を説明するModelSignaturesが推論時にLangchainモデルのアーティファクトとともに収集およびログに記録されます。
    log_models=True, # Trueの場合、LangchainモデルがMLflowモデルアーティファクトとしてログに記録されます
    log_inputs_outputs=True, # Trueの場合、推論データと結果が単一のDataFrameに結合され、推論データと結果がアーティファクトとしてMLflow Trackingにログに記録されます
    registered_model_name="lc_model",
)

# エージェント制御のための言語モデルをロードする
llm = OpenAI(temperature=0)

# 使用するツールをロードする
toolNames = ["serpapi"]
tools = load_tools(toolNames, llm=llm)

# ツール、言語モデル、そして使用したいエージェントのタイプを指定してエージェントを初期化します。
agent = initialize_agent(tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True)
# MLflowの実行内にエージェントをログに記録する
with mlflow.start_run():
    # エージェントツールをログに記録する
    mlflow.log_param("tool", toolNames) 
    # モデルを保存する
    logged_model = mlflow.langchain.log_model(agent, "langchain_model")

ex1.ipynd

# ログされたエージェントモデルをロードする
loaded_model = mlflow.pyfunc.load_model(logged_model.model_uri)

# ロードされたモデルを使用して推論結果を生成する
question = "明日の東京都の天気を教えてほしい"
answer = loaded_model.predict([{"input": question}])

# ロードされたモデルを使用して推論結果を生成する
question = "明日の北海道の天気を教えてほしい"
answer = loaded_model.predict([{"input": question}])

# ロードされたモデルを使用して推論結果を生成する
question = "明日の沖縄県の天気を教えてほしい"
answer = loaded_model.predict([{"input": question}])

結果：MLFLOWへの出力を確認

TOPページ

実行した結果が１件連携されていることを確認。オートログで記録できる項目が、metrixに表示されている

詳細画面：概要

カスタムで設定したパラメータ値がParametersに連携されていることを確認

詳細画面：アーティファクト

inference_inputs_outputs.jsonに実験した結果（東京・沖縄・北海道の天気についての推論結果）が格納されていることを確認

また、使用したトークン値などもメタデータとして自動ロギングされている

さらに、記録したモデルをロードして追加の推論を実行してみる

ex1.ipynd

# 予測のためにログされたエージェントモデルをロードする(RUIDを直で指定してみる)
loaded_model = mlflow.pyfunc.load_model("runs:/e56ebe853ee7446cbc46ac7a4ef54d73/langchain_model")

# ロードされたモデルを使用して推論結果を生成する
question = "明日の愛知県の天気を教えてほしい"
answer = loaded_model.predict([{"input": question}])

利用したモデルの結果に実験した推論結果が追記できたことを確認

※推論結果が少し違和感あるが、記事の観点でないので放置（２度目以降はちゃんとした結果を返してきたので、一時的な実行環境の問題かもです）。

結果

簡単な動作確認のみですが、MLFLOWとlangchainの連携が成功しました。

langchainの仕様に関する理解が不十分であり、連携されたデータの全体を把握できていませんが、必要な情報は適切に連携と管理がされていると考えています。

MLFLOWサーバをチーム内で共有すること、実験結果・内容の共有や他のメンバーがチューニングしたモデルでの検証などが容易に可能になると考えています。

簡単に実装できるため、コストをかけずに手軽に実験管理ができます。有効活用すれば効果的に運用できると思います。