DatabricksにおけるMLflowのLLMサポートのウォークスルー

この辺りを読んで、(多分)最近のMLflowをキャッチアップできたような気がします。

これらのまとめの意味を込めて、LLMOpsの文脈でMLflowをどのように活用できるのかをウォークスルーします。

全体像

MLflowのサイトでまとめられているように、LLM、LLMを活用したアプリケーション開発の複雑性を低減するために、MLflowから様々なサポートが提供されています。

LLM（大規模言語モデル）は、自然言語理解からコード生成まで幅広い能力を提供し、機械学習分野の基盤となっています。しかし、LLMの完全な潜在能力を引き出すには、複数のプロバイダーとのインターフェースから特定のモデルのファインチューニングまで、複雑なプロセスが伴います。

このような複雑さは、LLMの機能をアプリケーションに統合しようとする開発者やデータサイエンティストにとってボトルネックになりがちです。

MLflowのLLMサポートは、エンドユーザーを念頭に置いて設計された一連の機能とツールを導入することで、これらの課題を軽減することを目指しています。

ここ数年でのMLflowにおけるLLMサポートのアップデートは目覚ましいものがあります。以下に示すように、LLM(とそれを活用したアプリケーション)の構築、テスト、デプロイ、監視における様々なサポートを提供するようになっています。

• 実験トラッキング: MLflowではもともと提供されていた機能ですが、LLMサポートに際してチェーンを記録できるようになっています。
• トレース: LangSmithのように、生成AIアプリの動作をトレースすることができます。
• モデルパッケージ: OpenAI、Langchainなど生成AIに特化したフレーバーがどんどん追加されています。さらには、モデルをコードとして記録するModels From Codeがサポートされました。
• 評価: 生成AIの回答の正しさやLLM as a judgeによる評価がサポートされるようになりました。
• モデルサービング: REST APIからモデルを呼び出せるようにモデルサービングを行います。モデルサービングに関しては、Databricksのモデルサービングを活用することもできます。
• AI Gateway: 複数のLLMへのアクセスを一元管理できます。こちらもDatabricksのモデルサービングで代替することができます。
• プロンプトエンジニアリングUI: MLflowランの画面からプロンプトエンジニアリングを行えるようになりました。

ここでは、開発段階で頻繁に使うことになるであろう以下の機能にフォーカスしてウォークスルーします。

• トレース: LLMや生成AIアプリの開発の過程でチェーンの挙動をトレース。
• 実験トラッキング: メタデータと共にチェーンを記録。
• 評価: 開発したチェーンを評価。

準備

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

%pip install -U openai mlflow langchain langchain-openai langchain_community textstat
%restart_python

OpenAIのAPIキーを設定します。

import os
import openai
import pandas as pd
import mlflow

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

トレース

MLflowトレーシングは、LLM呼び出しやドキュメント取得、データクエリ、ツール呼び出しなどの重要な詳細をキャプチャすることで、GenAIアプリケーションにおけるLLMの可観測性を確保し、アプリケーションの内部動作を監視し深い洞察を得ることができます。

• デバッグ: トレースは抽象レイヤーの下で何が起こっているかの詳細なビューを提供し、LLMアプリケーションの問題を迅速に特定して解決するのに役立ちます。Jupyter Notebookとのネイティブ統合により、MLflowはシームレスなデバッグ体験を提供します。
• 品質の検査: MLflow Evaluationでモデルやエージェントを評価した後、自動生成されたトレースを分析してLLMの動作を理解し、情報に基づいた意思決定を行うことができます。
• 本番監視: トレースは本番環境でのLLMのパフォーマンスを監視するために不可欠です。MLflowトレーシングは、本番環境でのボトルネックやパフォーマンスの問題を特定し、是正措置を講じてアプリケーションを継続的に最適化することを可能にします。
• OpenTelemetry: MLflowトレースは、業界標準の可観測性フレームワークであるOpenTelemetryと互換性があり、Prometheus、Grafana、Jaegerなどの人気のある可観測性ツールと統合して高度な監視と分析を行うことができます。
• フレームワークサポート: MLflowトレーシングは、OpenAI、LangChain、LlamaIndex、DSPy、Anthropic、Amazon Bedrockなど、15以上のGenAIライブラリをサポートしています。これらのライブラリで構築されたLLMアプリケーションのトレースを開始するには、コードに1行追加するだけです。

from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_openai import ChatOpenAI

# LangChainの自動ログを有効にすると、トレースログが有効になります。
mlflow.langchain.autolog()

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7, max_tokens=1000)

prompt_template = PromptTemplate.from_template(
    "MLflowの専門家として答えてください。"
    "質問に対して高校生でもわかるように簡潔に説明してください。質問は: {question}"
)

# プロンプトテンプレート、LLM、出力パーサーから構成されるチェーン
chain = prompt_template | llm | StrOutputParser()

# チェーンの呼び出しをテストしましょう
chain.invoke(
    {
        "question": "MLflowとは何ですか？どのような時に使うものですか？",
    }
)

トレースからチェーンの挙動を確認できます。

実験トラッキング

実験トラッキングはMLflowプラットフォームの基本的な機能であり、モデル、プロンプト、トレース、メトリクスなど、LLM/GenAIプロジェクト中の多くのアセットやアーティファクトを一箇所に整理することができます。

• 単一の情報源: MLflowは、モデル、プロンプト、トレースなど、すべてのLLMアセットを保存および管理するための集中化された場所として機能し、プロジェクトの単一の情報源を確保します。
• より良いコラボレーション: 実験トラッキングにより、チームメンバーと作業を共有でき、結果を再現しフィードバックを提供することができます。
• リネージ追跡: MLflowは、プロジェクトのさまざまな段階を通じてLLMアセットのリネージをキャプチャし、モデル、プロンプト、その他のアーティファクトのリネージを追跡できるようにします。
• 比較分析: MLflowは、異なるバージョンのLLMモデル、プロンプト、その他のアセットを比較することを可能にし、情報に基づいた意思決定を行い、LLMアプリケーションの品質を向上させるのに役立ちます。
• ガバナンスとコンプライアンス: 組織内の機械学習アセットの中央リポジトリを使用することで、MLflowはガバナンスとコンプライアンス基準を維持し、重要なアセットが認可されたユーザーのみがアクセスできるようにします。

以降ではModels from codeアプローチを用いてチェーンをMLflowに記録します。

Models from code機能は、シリアライズされたモデルの重みに依存しないカスタムモデルや(LangChainやLlamaIndexのような)特定のフレーバー実装の両方に対する定義、格納、ロードのプロセスを包括的に刷新します。

これらのモデルにおけるレガシーなシリアライゼーションと、Models from Codeアプローチの主な違いは、シリアライゼーションの過程でどのようにモデルが表現されるかです。

レガシーなアプローチでは、cloudpickle(カスタムのpyfuncとLangChain)や、背後のパッケージにおける全ての機能に対する網羅性が不完全なカスタムのシリアライザ(LlamaIndexの場合)のどちらかを用いてモデルオブジェクトに対してシリアライゼーションが行われます。カスタムpyfuncでは、オブジェクトインスタンスのシリアライズにcloudpickleを用いることで、ロードの際にオブジェクトを再構成するために使用されるバイナリーファイルが作成されます。

Models From Codeでは、サポートされるモデルタイプにおいては、カスタムpyfuncやフレーバーのインタフェース(LangChainの場合、スクリプト内に直接LCELチェーンを定義、マークすることができます)の定義とともにシンプルなスクリプトが保存されます。

まず、チェーンを実装する.pyファイルを保存します。

%%writefile "./chain_1.py"

import mlflow
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7, max_tokens=1000)

# 標準的なシステムプロンプト
prompt_template_1 = PromptTemplate.from_template(
    "あなたの仕事はMLflowに関する質問に答えることです。質問は: {question}"
)

# プロンプトテンプレート、LLM、出力パーサーから構成されるチェーン
chain_1 = prompt_template_1 | llm | StrOutputParser()

mlflow.models.set_model(chain_1)

Overwriting ./chain_1.py

%%writefile "./chain_2.py"

import mlflow
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7, max_tokens=1000)

# 関係のない話題に対する回答を拒否するシステムプロンプト
prompt_template_2 = PromptTemplate.from_template(
    "あなたの仕事はMLflowに関する質問に答えることです。MLflowに関する質問があった場合、"
    "それに答えてください。コード例を含めるようにしてください。質問がMLflowに関連していない場合、"
    "答えを拒否し、その質問が無関係であることを伝えてください。質問は: {question}"
)

# プロンプトテンプレート、LLM、出力パーサーから構成されるチェーン
chain_2 = prompt_template_2 | llm | StrOutputParser()

mlflow.models.set_model(chain_2)

Overwriting ./chain_2.py

chain_1.pyをコードとしてMLflowにロギングします。

chain_path = "./chain_1.py"
input_example = "MLflowとは何ですか？"

with mlflow.start_run():
    mlflow.log_param("chain name", "chain_1")
    mlflow.log_param("chain type", "通常のMLflowアシスタント")

    info = mlflow.langchain.log_model(
        lc_model=chain_path,  # チェーン定義とset_model呼び出しを含むスクリプトとしてモデルを定義
        artifact_path="chain",
        input_example=input_example,
    )

chain_path = "./chain_2.py"
input_example = "MLflowとは何ですか？"

with mlflow.start_run():
    mlflow.log_param("chain_name", "chain_1")
    mlflow.log_param("chain_type", "MLflowに関係のない質問への回答を拒否するアシスタント")

    info = mlflow.langchain.log_model(
        lc_model=chain_path,  # チェーン定義とset_model呼び出しを含むスクリプトとしてモデルを定義
        artifact_path="chain",
        input_example=input_example,
    )

二つのモデルとして記録されました。

記録したモデルをロードしてみます。以下のコードではchain_2.pyのチェーンがロードされます。

loaded_chain = mlflow.langchain.load_model(model_uri=info.model_uri)

question = "MLflowとは何ですか？"

response = loaded_chain.invoke(question)
print(response)

loaded_chain

PromptTemplate(input_variables=['question'], input_types={}, partial_variables={}, template='あなたの仕事はMLflowに関する質問に答えることです。MLflowに関する質問があった場合、それに答えてください。コード例を含めるようにしてください。質問がMLflowに関連していない場合、答えを拒否し、その質問が無関係であることを伝えてください。質問は: {question}')
| ChatOpenAI(client=<openai.resources.chat.completions.completions.Completions object at 0x7ff0cda93260>, async_client=<openai.resources.chat.completions.completions.AsyncCompletions object at 0x7ff0cd72eb10>, root_client=<openai.OpenAI object at 0x7ff0cd72f3e0>, root_async_client=<openai.AsyncOpenAI object at 0x7ff0cda91280>, model_name='gpt-4o-mini', temperature=0.7, model_kwargs={}, openai_api_key=SecretStr('**********'), max_tokens=1000)
| StrOutputParser()

評価

MLflowのLLM評価は、基盤モデル、プロンプト、および複合AIシステムを比較するための簡素化されたアプローチを提供し、評価プロセスを簡素化するように設計されています。

• 簡素化された評価: MLflowのmlflow.evaluate() APIと組み込みのハーネスを活用して、GenAIモデルを簡単にテストできます。
• LLM-as-a-Judge: MLflowは柔軟な設定が可能なLLM-as-a-Judgeメトリクスを提供し、評価プロセスに複雑な評価基準を組み込んで一括評価を行うことができます。
• カスタマイズ可能なメトリクス: 提供されるメトリクスに加えて、MLflowはプラグインスタイルのカスタムスコアリングをサポートし、評価の柔軟性を高めます。
• 比較分析: 基盤モデル、プロバイダー、プロンプトを簡単に比較して、情報に基づいた意思決定を行うことができます。
• 深い洞察: 評価実行のためにトレースが自動的に生成され、評価結果に関する深い洞察を得て、品質問題を自信を持って修正できます。

from mlflow.models import infer_signature

# チェーンの記録と評価を実施
def build_and_evaluate_model_with_chain(chain_path, chain_name, chain_type):
    mlflow.start_run()
    mlflow.log_param("chain_name", chain_name)
    mlflow.log_param("chain_type", chain_type)

    # 入力サンプル
    input_example = "MLflowとは何ですか？"

    # Langchainを使用して質問応答モデルを作成します。モデルをMLflow Trackingにログします
    logged_model = mlflow.langchain.log_model(
        lc_model=chain_path,  # チェーン定義とset_model呼び出しを含むスクリプトとしてモデルを定義
        artifact_path="chain",
        input_example=input_example,
    )

    # いくつかの例題でモデルを評価します
    questions = pd.DataFrame(
        {
            "question": [
                "MLflowでランを作成するにはどうすればよいですか？",
                "MLflowでモデルをログするにはどうすればよいですか？",
                "フランスの首都はどこですか？",
            ]
        }
    )

    # 評価データセットを用いてチェーンを評価します
    mlflow.evaluate(
        model=logged_model.model_uri,
        model_type="question-answering",
        data=questions,
    )
    mlflow.end_run()

print("chain_1を構築および評価しています:")
build_and_evaluate_model_with_chain("./chain_1.py", "chain_1", "通常のMLflowアシスタント")

print("chain_2を構築および評価しています:")
build_and_evaluate_model_with_chain("./chain_2.py", "chain_2", "MLflowに関係のない質問への回答を拒否するアシスタント")

すると、今回はモデルに加えて評価結果も記録されます。

エクスペリメントの画面ではより詳細に評価結果を確認することができます。

# 評価結果をロードして確認します
results: pd.DataFrame = mlflow.load_table(
    "eval_results_table.json", extra_columns=["run_id", "params.chain_type"]
)
results_grouped_by_question = results.sort_values(by="question")
print("評価結果:")
display(results_grouped_by_question[["run_id", "params.chain_type", "question", "outputs"]])

run_id	params.chain_type	question	outputs
82730be0f6474a40b276b56f5fc5844f	MLflowに関係のない質問への回答を拒否するアシスタント	MLflowでモデルをログするにはどうすればよいですか？	MLflowでモデルをログするには、MLflowのmlflowライブラリを使ってモデルのトレーニング後にモデルを保存し、ログすることができます。以下は、MLflowを使用してモデルをログする基本的な手順の例です。 まず、MLflowをインストールしておく必要があります。次のコマンドでインストールできます。 bash<br>pip install mlflow<br> 次に、以下のようにモデルをトレーニングし、ログすることができます。ここでは、Scikit-learnの線形回帰モデルを例に示します。 python<br>import mlflow<br>import mlflow.sklearn<br>from sklearn.linear_model import LinearRegression<br>from sklearn.datasets import make_regression<br>from sklearn.model_selection import train_test_split<br><br># データセットの作成<br>X, y = make_regression(n_samples=100, n_features=1, noise=10)<br>X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)<br><br># モデルのトレーニング<br>model = LinearRegression()<br>model.fit(X_train, y_train)<br><br># モデルのログ<br>with mlflow.start_run():<br> mlflow.sklearn.log_model(model, "model")<br> mlflow.log_param("fit_intercept", model.fit_intercept)<br> mlflow.log_param("normalize", model.normalize)<br><br># モデルをMLflowサーバーに保存<br>mlflow.log_artifact("path/to/your/artifact") # 任意でアーティファクトをログする場合<br> このコードでは、まずデータセットを作成し、線形回帰モデルをトレーニングします。その後、mlflow.start_run()を使用して新しい実行を開始し、mlflow.sklearn.log_modelを使ってモデルをログします。また、モデルのパラメータもmlflow.log_paramでログしています。 これで、MLflowのUIからモデルの実行やログされたパラメータ、メトリクス、モデルアーティファクトを確認することができます。
03c206e2aafa48e4a6ee7fd54b4f9978	通常のMLflowアシスタント	MLflowでモデルをログするにはどうすればよいですか？	MLflowでモデルをログするには、以下の手順に従います。MLflowは、機械学習モデルのトラッキング、パッケージング、デプロイを簡単に行えるオープンソースのプラットフォームです。 1. MLflowのインポート: まず、MLflowをインポートします。必要に応じて、他のライブラリ（例: scikit-learn、TensorFlow、PyTorchなど）もインポートします。 python<br> import mlflow<br> import mlflow.sklearn # 例として、scikit-learnを使用する場合<br> 2. MLflowのトラッキングサーバーの設定: トラッキングサーバーを使用している場合は、適切なURIを設定します。 python<br> mlflow.set_tracking_uri("http://your-tracking-server:5000")<br> 3. 実験の開始: 新しい実験を作成するか、既存の実験に参加します。 python<br> mlflow.create_experiment("your_experiment_name") # 新しい実験を作成<br> mlflow.start_run() # 新しいランを開始<br> 4. モデルのトレーニング: モデルをトレーニングします。ここでは例としてscikit-learnのモデルを使用します。 python<br> from sklearn.ensemble import RandomForestClassifier<br> from sklearn.datasets import load_iris<br> from sklearn.model_selection import train_test_split<br><br> # データの準備<br> data = load_iris()<br> X_train, X_test, y_train, y_test = train_test_split(data.data, data.target, test_size=0.2)<br><br> # モデルのトレーニング<br> model = RandomForestClassifier()<br> model.fit(X_train, y_train)<br> 5. モデルをログする: トレーニングしたモデルをMLflowにログします。 python<br> mlflow.sklearn.log_model(model, "model")<br> 6. パラメータやメトリクスのログ（オプション）: モデルのハイパーパラメータや評価メトリクスをログすることもできます。 python<br> mlflow.log_param("n_estimators", model.n_estimators)<br> mlflow.log_metric("accuracy", model.score(X_test, y_test))<br> 7. ランの終了: 最後に、ランを終了します。 python<br> mlflow.end_run()<br> これで、MLflowにモデルをログする準備が整いました。MLflow UIを使用して、ログされたモデルやメトリクスを確認することができます。MLflowの詳細な使い方については、公式ドキュメントを参照することをお勧めします。
82730be0f6474a40b276b56f5fc5844f	MLflowに関係のない質問への回答を拒否するアシスタント	MLflowでランを作成するにはどうすればよいですか？	MLflowでランを作成するには、まずMLflowをインポートし、mlflow.start_run()を使用して新しいランを開始します。その後、モデルのパラメータやメトリクスをログすることができます。以下は、MLflowでランを作成する基本的なコード例です。 python<br>import mlflow<br><br># MLflowのトラッキングサーバーに接続（必要に応じて）<br>mlflow.set_tracking_uri("http://localhost:5000")<br><br># ランを開始<br>with mlflow.start_run() as run:<br> # モデルのパラメータをログ<br> mlflow.log_param("param1", 5)<br> mlflow.log_param("param2", "value2")<br> <br> # モデルのメトリクスをログ<br> mlflow.log_metric("metric1", 0.85)<br> mlflow.log_metric("metric2", 0.95)<br> <br> # モデルを保存する場合の例（ここでは仮にモデルオブジェクトがあるとします）<br> # mlflow.sklearn.log_model(model, "model")<br><br> # ランIDを取得<br> run_id = run.info.run_id<br> print(f"Run ID: {run_id}")<br> このコードでは、MLflowを使って新しいランを作成し、いくつかのパラメータとメトリクスをログしています。ランが終了すると、自動的にMLflowに保存されます。必要に応じて、モデルも保存することができます。
03c206e2aafa48e4a6ee7fd54b4f9978	通常のMLflowアシスタント	MLflowでランを作成するにはどうすればよいですか？	MLflowでランを作成するには、以下の手順を実行します。MLflowは実験のトラッキング、モデルの管理、プロジェクトのパッケージ化などをサポートするプラットフォームです。 1. MLflowのインストール: まず、MLflowをインストールする必要があります。以下のコマンドを使用してインストールできます。 bash<br> pip install mlflow<br> 2. MLflowのインポート: PythonスクリプトやノートブックでMLflowを使用するためにインポートします。 python<br> import mlflow<br> 3. ランの開始: mlflow.start_run()を使用して新しいランを開始します。この関数は、コンテキストマネージャとして使用することができ、ブロック内のコードに関連付けられたランを自動的に終了します。 python<br> with mlflow.start_run():<br> # モデルのトレーニングや評価のコードをここに書きます<br> ...<br> # メトリックやパラメータをログします<br> mlflow.log_param("param_name", param_value)<br> mlflow.log_metric("metric_name", metric_value)<br> 4. パラメータやメトリックのログ: ランの中で、モデルのパラメータや評価メトリックをログすることができます。 python<br> mlflow.log_param("learning_rate", 0.01)<br> mlflow.log_metric("accuracy", 0.95)<br> 5. モデルの保存（オプション）: モデルをトレーニングした後、MLflowを使用してモデルを保存することもできます。 python<br> mlflow.sklearn.log_model(model, "model")<br> 6. ランの終了: コンテキストマネージャを使用している場合、ランはブロックの終わりで自動的に終了します。手動で終了する場合は、以下のようにします。 python<br> mlflow.end_run()<br> これで、MLflowを使用してランを作成し、パラメータやメトリックをログし、必要に応じてモデルを保存することができます。ランの結果はMLflowのUIで確認することができ、実験の管理が容易になります。
82730be0f6474a40b276b56f5fc5844f	MLflowに関係のない質問への回答を拒否するアシスタント	フランスの首都はどこですか？	その質問はMLflowに関連していないため、お答えできません。MLflowに関する質問があれば、ぜひお聞きください。
03c206e2aafa48e4a6ee7fd54b4f9978	通常のMLflowアシスタント	フランスの首都はどこですか？	フランスの首都はパリです。ただし、MLflowに関する質問があれば、お気軽にどうぞ！

# 新しい質問で最良のモデルをスコアリングします
new_question = "MLflow Model Registryでモデルバージョンを作成するにはどうすればよいですか？"
best_model = mlflow.pyfunc.load_model(f"runs:/{mlflow.last_active_run().info.run_id}/chain")
response = best_model.predict(new_question)
print(f"回答: {response}")

回答: ['MLflow Model Registryでモデルバージョンを作成するには、まずモデルをMLflowにログし、その後モデルを登録する必要があります。以下は、基本的な手順とコード例です。\n\n1. **モデルのトレーニング**: まず、モデルをトレーニングします。\n2. **モデルをログに記録**: 学習したモデルをMLflowにログします。\n3. **モデルを登録**: ログしたモデルをModel Registryに登録します。\n\n以下は、Pythonでの具体的なコード例です。\n\n```python\nimport mlflow\nimport mlflow.sklearn\nfrom sklearn.ensemble import RandomForestClassifier\nfrom sklearn.datasets import load_iris\nfrom sklearn.model_selection import train_test_split\n\n# データのロード\ndata = load_iris()\nX_train, X_test, y_train, y_test = train_test_split(data.data, data.target, test_size=0.2, random_state=42)\n\n# モデルのトレーニング\nmodel = RandomForestClassifier()\nmodel.fit(X_train, y_train)\n\n# MLflowの追跡の開始\nwith mlflow.start_run() as run:\n    # モデルをログに記録\n    mlflow.sklearn.log_model(model, "model")\n    \n    # モデルを登録\n    model_uri = f"runs:/{run.info.run_id}/model"\n    mlflow.register_model(model_uri, "IrisModel")\n\n# これで、モデルがModel Registryに登録されました。\n```\n\nこのコードでは、Irisデータセットを使用してランダムフォレストモデルをトレーニングし、MLflowを使用してモデルをログし、Model Registryに登録しています。登録後、モデルはModel Registry内でバージョン管理されるようになります。\n\nモデルのバージョンを確認するには、MLflow UIを開き、"Models"タブを選択することで、登録されたモデルの情報を確認できます。']

はじめてのDatabricks

はじめてのDatabricks

Databricks無料トライアル

Databricks無料トライアル