MLflow 2.3のLangChainのサポートを試す

こちらのLangChainのサポートを試します。[LangChain](https://python.langchain.com/en/latest/index.html#)触るの初めてです。アイコンが良い🦜🔗

https://qiita.com/taka_yayoi/items/431fa69430c5c6a5e741#langchain%E3%81%AE%E3%82%B5%E3%83%9D%E3%83%BC%E3%83%88

langchainとmlflow 2.3をインストールします。

```
%pip install langchain mlflow==2.3
```

# モデルの記録

`LLMChain`の`HuggingFaceHub`でリポジトリにアクセスする際には、環境変数`HUGGINGFACEHUB_API_TOKEN`にHugging FaceのAPIトークンを設定します。
![Screenshot 2023-04-21 at 15.39.21.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/489d2124-3fa5-6403-b221-7d3438855a68.png)

また、モデルをロードする際にもAPIトークンが必要になるので、Databricksシークレットに格納の上、クラスターで環境変数が設定されるようにします。

```
databricks secrets put --scope demo-token-takaaki.yayoi --key huggingface_api_key
```

クラスターの設定でシークレットを参照するようにします。
![Screenshot 2023-04-21 at 15.43.22.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/01e842a3-469a-33e2-9006-5a410b19b935.png)

ここでは`google/flan-t5-small`を使用しています。

```py:Python
import os
import mlflow

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-small", # モデルが大きいとロード時にエラーになります
        model_kwargs={"temperature":0, "max_length":64}
    ),
)

# チェーンの記録
mlflow.langchain.log_model(
    lc_model=llm_chain,
    artifact_path="model",
    registered_model_name="english-to-french-chain-gpt-3-5-turbo-1"
)
```

# モデルのロード

モデルレジストリにも登録されているので、`model:`のURIでモデルを呼び出すことができます。

```py:Python
import mlflow.pyfunc

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

Spark UDFとしてモデルをロードしているので、大量データに対してもLLMの処理を並列実行することができます。

翻訳対象のデータを準備します。

```py:Python
import pandas as pd

english_df = pd.DataFrame(
    {
        "content": [
            "What is MLflow?",
            "What are the key components of MLflow?",
            "How does MLflow enable reproducibility?",
            "What is MLflow tracking and how does it help?",
            "How can you compare different ML models using MLflow?",
            "How can you use MLflow to deploy ML models?",
            "What are the integrations of MLflow with popular ML libraries?",
            "How can you use MLflow to automate ML workflows?",
            "What security and compliance features does MLflow offer?",
            "Where does MLflow stand in the ML ecosystem?",
        ],
    }
)

englsh_sdf = spark.createDataFrame(english_df)
display(englsh_sdf)
```
![Screenshot 2023-04-21 at 16.07.59.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/c8b6d9fc-7999-43cc-3da1-1a267f331e7c.png)

# LLMの呼び出し

通常のSpark UDFと同じようにSparkデータフレームにLLMを適用することができます。

```py:Python
french_translated_df = englsh_sdf.withColumn(
    "french_text",
    english_to_french_udf("content")
)   
display(french_translated_df)
```

残念ですが私はフランス語分かりませんが、うまく動いているようです。
![Screenshot 2023-04-21 at 16.09.13.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/dd64d530-62a1-4bf1-01b8-2a9b524eb133.png)

これまでの機械学習モデルのトラッキングと若干毛色が違いますが、LLMを試行錯誤して色々行うのが便利になりそうです。LangChainをもっと勉強しよう。

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

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


### Databricks無料トライアル

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