MLflow 2.3のご紹介：ネイティブLLMのサポートと新機能による強化

[Introducing MLflow 2\.3: Enhanced with Native LLM Support and New Features \- The Databricks Blog](https://www.databricks.com/blog/2023/04/18/introducing-mlflow-23-enhanced-native-llm-support-and-new-features.html)の翻訳です。

:::note warn
本書は抄訳であり内容の正確性を保証するものではありません。正確な内容に関しては原文を参照ください。
:::

月間1300万以上のダウンロードによって、MLflowはエンドツーエンドのMLOpsプラットフォームとしての最高の地位を確立し、いかなるサイズのチームが、バッチ、リアルタイム推論の両方において、モデルを追跡、共有、パッケージ、デプロイすることを支援しています。プロダクションの機械学習アプリケーションをドライブするために数千の企業が日々MLflowを活用しており、業界、学術機関からの500以上のコントリビューターによるコミュニティによってアクティブに開発されています。

本日、大規模言語モデル(LLM)を管理、デプロイできる能力を拡張する革新的な機能が搭載された、このオープンソース機械学習プラットフォームの最新のアップデートであるMLflow 2.3を公開できることを嬉しく思っています。この強化されたLLMサポートは以下を通じて提供されます:

- 3つの新たなモデルフレーバー: [Hugging Face Transformers](https://huggingface.co/docs/transformers/index)、[OpenAI](https://openai.com/product)関数、[LangChain](https://python.langchain.com/en/latest/index.html)
- モデルファイルのマルチパートのダウンロード、アップロードを通じたクラウドサービスとのモデルのダウンロード、アップロード速度の劇的な改善

# Hugging Faceトランスフォーマーのサポート

「MLflowにいおけるHugging Faceトランスフォーマーのネイティブサポートによって、AIの最大のコミュニティ、オープンプラットフォームであるHugging Face Hubで利用できる[170,000以上の無料で公開されている機械学習モデル](https://huggingface.co/models)を容易に取り扱えるようになります。」- [Jeff Boudier](mailto:jeff@huggingface.co), Product Director, Hugging Face

この新たなトランスフォーマーのフレーバーによって、MLflowトラッキングサービスにトランスフォーマーパイプライン、モデルのネイティブのインテグレーションと処理コンポーネントが提供されます。この新たなフレーバーによって、共通MLflowトラッキングインタフェースを通じて、完全に設定されたトランスフォーマー[パイプライン](https://huggingface.co/docs/transformers/pipeline_tutorial)や[Dolly](https://huggingface.co/databricks/dolly-v2-12b)を含むベースモデルを記録することが可能になります。これらの記録されたアーティファクトは、コンポーネントのコレクション、パイプライン、[pyfunc](https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html)を通じてネイティブにロードすることができます。

トランスフォーマーフレーバーを用いてコンポーネントやパイプラインを記録する際、数多くの検証が自動で実行され、保存するパイプラインやモデルがデプロイ可能なインタフェースと互換性を持つことが保障されます。この検証に加え、重要なモデルカードデータが自動で取得され、モデルやパイプラインのアーティファクトに追加されます。これらのモデルやパイプラインの使いやすさを改善するために、デプロイメントプロセスをシンプルにするために、モデルシグネチャインタフェース機能が追加されます。

Hugging Face HubにホスティングされているDollyのようなオープンソース大規模言語モデルとのインテグレーションは以下のようにシンプルです:

```py:Python
import mlflow
import transformers

architecture = "databricks/dolly-v2-3b"

dolly = transformers.pipeline(model=architecture, trust_remote_code=True)

with mlflow.start_run():
    model_info = mlflow.transformers.log_model(
        transformers_model=dolly,
        artifact_path="dolly3b",
        input_example="Hello, Dolly!",
    )

loaded_dolly = mlflow.transformers.load_model(
model_info.model_uri, 
max_new_tokens=250,
)

>>> [{'generated_text': 'Please note that MLflow is not a "ML IDE". That means that MLflow is not your only tool for exploring and building ML pipelines. However, MLflow provides a pipeline engine that you can use to automatically train and deploy models for serving in production in the cloud or on-premises. To get started, we recommend checking out the MLflow quick start guide. If you are a non-technical user, we also recommend the Getting Started page to understand the end-to-end ML experience.'}]
```

この新たなHugging Faceトランスフォーマーとのインテグレーションによって、以下のように軽量チャットボットインタフェースとしてのpyfuncバージョンのモデルを活用することすら可能になります。

```py:Python
import transformers
import mlflow

chat_pipeline = transformers.pipeline(model="microsoft/DialoGPT-medium")

with mlflow.start_run():
  model_info = mlflow.transformers.log_model(
    transformers_model=chat_pipeline,
    artifact_path="chatbot",
    input_example="Hi there!"
  )

# Load as interactive pyfunc

chatbot = mlflow.pyfunc.load_model(model_info.model_uri)
```

このMLflowのトランスフォーマーフレーバーは、自動シグネチャスキーマ検知と、パイプライン固有の入力フォーマットの指定をサポートしています。

トランスフォーマーパイプラインからロードされたpyfuncモデルの活用は、以下に示すように背後のパイプラインのインタフェースを保持することを目的としています。

```py:Python
chatbot.predict("What is the best way to get to Antarctica?")

>>> 'I think you can get there by boat'

chatbot.predict("What kind of boat should I use?")

>>> 'A boat that can go to Antarctica.'
```

トランスフォーマーパッケージによってサポートされているパイプラインタイプのそれぞれにおいて、保存されたモデルやパイプラインをあとで参照したり、サービングする際に、正確な要件、コンポーネントのバージョン、参照情報を利用できるようにメタデータが収集されます。明示的に宣言された入力シグネチャなしでも、記録の際に指定された代表的な入力サンプルに基づいてシグネチャが推定されます。
![](https://cms.databricks.com/sites/default/files/inline-images/db-602-blog-img-1.png)
*MLflowが自動でモデルからメタデータを抽出します*

さらに、MLflowトランスフォーマーフレーバーは、モデルやパイプラインを記録する際に、Hugging Face Hubからモデルカードの状態を自動で取得します。この機能によって、一般的な参照や監査目的の両方において、背後のモデル情報の状態に対するポイントインタイムの参照が可能になります。
![](https://cms.databricks.com/sites/default/files/inline-images/db-602-blog-img-2.png)
*また、MLflowはHugging Faceからモデルカードを取得します*

このトランスフォーマーフレーバーの追加によって、広く活用されている人気のパッケージがMLflowでファーストクラスのサポートを受けるようになります。MLflowを用いて再トレーニング、記録されたパイプラインを容易に[Hugging Face Repository](https://huggingface.co/docs/transformers/model_sharing)に送信し直すことができるので、他の人たちが複雑なテキストベースの問題を解くために、最新のモデルアーキテクチャを活用し、メリットを享受できるようになります。

# 関数ベースのフレーバーによるOpenAI APIのサポート

[OpenAI Pythonライブラリ](https://github.com/openai/openai-python)によって、Python言語で記述されたアプリケーションからOpenAI APIに簡単にアクセスできるようになります。これには、OpenAI APIリソースへのマッピングを行う一連の事前定義積みクラスが含まれています。これら提供されるクラスを用いることで、さまざまなバージョンのモデルやOpenAI APIエンドポイントに対して、動的に接続を初期化し、データを引き渡し、レスポンスを取得します。

MLflow OpenAIフレーバーは以下をサポートしています:

- 自動シグネチャスキーマ検知
- 推論を加速するための並列APIリクエスト
- レートリミットエラーのような一時的なエラーにおける自動APIリクエストリトライ

以下には`openai.ChatCompletion`モデルの記録、推論のためのロードのサンプルを示しています:

```py:Python
import mlflow
import openai

with mlflow.start_run();:
    model_info = mlflow.openai.log_model(
        model="gpt-3.5-turbo",
        task=openai.ChatCompletion,
        messages=[{"role": "system", "content": "You are an MLflow expert"}],
        artifact_path="model",
    )

model = mlflow.pyfunc.load_model(model_info.model_uri)

print(model.predict([{"role": "user", "content": "What is MLflow?"]))

# -> "MLflow is an open source ...."
```

MLflow 2.3では、カスタムモデルのロギングをシンプルにするために、モデルとしての関数の記録と、型アノテーションからの自動シグネチャ検知をサポートしています。以下には、関数型のOpenAIチャットコンプリーションモデルの記録の例を示しています。

```py:Python
from typing import List
import openai
import mlflow

# Define a functional model with type annotations

def chat_completion(inputs: List[str]) -> List[str]:
    # Model signature is automatically constructed from
    # type annotations. The signature for this model
    # would look like this:
    # ----------
    # signature:
    #   inputs: [{"type": "string"}]
    #   outputs: [{"type": "string"}]
    # ----------

    outputs = []

    for input in inputs:
        completion = openai.ChatCompletion.create(
            model="gpt-3.5-turbo"
            messages=[{"role": "user", "content": "<prompt>"}]
        )

        outputs.append(completion.choices[0].message.content)

    return outputs

# Log the model
mlflow.pyfunc.log_model(
    artifact_path="model",
    python_model=chat_completion,
    pip_requirements=["openai"],
)
```

MLflowでOpenAIフレーバーを活用することで、MLflowのトラッキング、デプロイメント機能を活用しつつも、OpenAIでホスティングされている事前定義済みモデルを完全に活用することが可能となります。

OpenAI APIキーは環境変数で管理し、ランのパラメーターやタグとして記録することは避けるようにしていください。Databricksでは、以下のように`openai_api_key`といったキー名とスコープでAPIキーを追加するために、[Databricksシークレット管理](https://qiita.com/taka_yayoi/items/7e4eae0786b94cb7a5de)を活用することができます。OpenAIフレーバーのモデルがエンドポイントでサービングされる際、MLflowは自動でDatabricksのシークレットストアからシークレットキーを取得します。

```
databricks secrets put --scope <scope-name> --key openai_api_key
```

シークレットのスコープ名は、環境変数`MLFLOW_OPENAI_SECRET_SCOPE`で指定することができます。

```py:Python
import os
import mlflow

os.environ["MLFLOW_OPENAI_SECRET_SCOPE"] = "<scope-name>"

# When the MLFLOW_OPENAI_SECRET_SCOPE variable is set, 
# `mlflow.openai.log_model` reads its value and saves it in `openai.yaml`

mlflow.openai.log_model(...)
```

# LangChainのサポート

MLflowの[LangChain](https://python.langchain.com/en/latest/index.html#)フレーバーは、質問応答システムやチャットボットのようなLLMベースのアプリケーションの構築やデプロイのプロセスをシンプルにします。LangChainの高度な機能と、MLflowの効率化された開発、デプロイサポートを活用することができます。

以下に、ネイティブのLangChainフレーバーとしてLLMChain(英語からフランス語への翻訳)をどのように記録し、Sparkを用いて英語テキストのバッチ翻訳を行う例を示します:

```py:Python
from langchain import PromptTemplate, HuggingFaceHub, LLMChain

template = """Translate everything you see after this into French:

{input}"""

prompt = PromptTemplate(template=template, input_variables=["input"])

llm_chain = LLMChain(
    prompt=prompt,
    llm=HuggingFaceHub(
        repo_id="google/flan-t5-xl",
        model_kwargs={"temperature":0, "max_length":64}
    ),
)

mlflow.langchain.log_model(
    model=llm_chain,
    artifact_path="model",
    registered_model_name="english-to-french-chain-gpt-3.5-turbo-1"
)
```

Spark UDFを用いた分散バッチ推論のためのLangChainモデルをロードします:

```py:Python
import mlflow.pyfunc

english_to_french_udf = mlflow.pyfunc.spark_udf(
    spark=spark,
    model_uri="models:/english-to-french-gpt-3.5-turbo-1/1",
    result_type="string"
)

french_translated_df = english_df.withColumn(
    "french_text",
    english_to_french_udf("english_text")
)   
```

これはLangChainフレーバーの初期リリースであり、モデルのロギングはLLMChainに限定されていることに注意してください。エージェントのロギングやその他のロギング昨日は開発中であり、今後のMLflowリリースで追加される予定です。

# MLflow 2.3を使い始める

すぐにでもMLflow 2.3を試していたくことをお勧めします！LLMをサポートする新機能を追加始めるには、シンプルに以下のコマンドを用いて[Python MLflow library](https://pypi.org/project/mlflow/)をインストールするだけです:

```
pip install mlflow==2.3
```

MLflow 2.3における新機能と改善点の完全なリストについては、[release changelog](https://github.com/mlflow/mlflow/blob/master/CHANGELOG.md)をご覧ください。MLflowを使い始める方法の詳細や、MLflow 2.3におけるLLMベースの機能の完全なドキュメントについては、[MLflow documentation](https://mlflow.org/docs/latest/index.html)をご覧ください。

### Databricksクイックスタートガイド

[Databricksクイックスタートガイド](https://www.amazon.co.jp/dp/B09V1YXFVQ/)


### Databricks無料トライアル

[Databricks無料トライアル](https://databricks.com/jp/try-databricks)