LLM 開発用の Azure Databricks 上の LangChain

この記事では、Azure Databricks での大規模言語モデル (LLM) の開発とデプロイを容易にする LangChain 統合について説明します。

これらの LangChain 統合により、次のことが可能になります。

• PySpark DataFrame ローダーを使用して、PySpark DataFrame からデータをシームレスに読み込みます。
• Spark DataFrame Agent または Databricks SQL Agent を使用して、自然言語を使用してデータに対話形式でクエリを実行します。
• Azure Databricks で提供されるモデルを LangChain の大規模言語モデル (LLM) としてラップします。

LangChain とは

LangChain は、大規模言語モデル (LLM) を利用するアプリケーションの作成を支援するように設計されたソフトウェア フレームワークです。 LangChain の強みは、その幅広い統合と機能にあります。 これには、API ラッパー、Web スクレイピング サブシステム、コード分析ツール、ドキュメント要約ツールなどが含まれます。 また、OpenAI、Anthropic、HuggingFace などの大規模な言語モデルを、さまざまなデータ ソースやタイプとともにすぐにサポートします。

LangChain は、LangChain のお客様が Azure Databricks 環境から直接 MLflow の堅牢なツールと実験追跡機能を活用できるようにする、実験的な MLflow フレーバーとして利用できます。 LangChain フレーバーの MLflow ドキュメントを参照してください。

必要条件

• Databricks Runtime 13.3 ML 以降。
• Databricks では、最新の更新を確実に入手できるように、最新バージョンの LangChain を pip インストールすることをお勧めします。
  • %pip install --upgrade langchain

PySpark DataFrame ローダーを使用してデータを読み込む

LangChain の PySpark DataFrame ローダーは、単一のメソッドで PySpark DataFrame からのデータのロードを簡素化します。

from langchain.document_loaders import PySparkDataFrameLoader

loader = PySparkDataFrameLoader(spark, wikipedia_dataframe, page_content_column="text")
documents = loader.load()

次のノートブックは、PySpark DataFrame ローダーを使用して、MLflow でログに記録される取得ベースのチャットボットを作成する例を示しています。これにより、モデルが mlflow.pyfunc.load_model() による推論用の汎用 Python 関数として解釈されるようになります。

Langchain ノートブックの PySpark DataFrame ローダーと MLflow

ノートブックを入手

Spark DataFrame エージェント

LangChain の Spark DataFrame エージェントを使用すると、質問応答用に最適化された Spark DataFrame との対話が可能になります。 LangChain の Spark DataFrame エージェント ドキュメントには、DataFrame で Spark DataFrame エージェントを作成および使用する方法の詳細な例が記載されています。

from langchain.agents import create_spark_dataframe_agent

df = spark.read.csv("/databricks-datasets/COVID/coronavirusdataset/Region.csv", header=True, inferSchema=True)
display(df)

agent = create_spark_dataframe_agent(llm=OpenAI(temperature=0), df=df, verbose=True)
...

次のノートブックでは、データに関する分析情報を得るために Spark DataFrame エージェントを作成および使用する方法を示します。

LangChain を使用して Spark DataFrame ノートブックと対話する

ノートブックを入手

Databricks SQL エージェント

Databricks SQL エージェントは、LangChain が提供する標準 SQL Database エージェントのバリアントであり、Spark DataFrame エージェントのより強力なバリアントとみなされます。

Databricks SQL エージェントを使用すると、Azure Databricks ユーザーは Unity Catalog 内の指定されたスキーマを操作し、データに関する分析情報を生成できます。

次の例では、データベース インスタンスが SQLDatabase.from_databricks(catalog="...", schema="...") コマンド内で作成され、エージェントと必要なツールがそれぞれ SQLDatabaseToolkit(db=db, llm=llm) と create_sql_agent(llm=llm, toolkit=toolkit, **kwargs) によって作成されます。

from langchain.agents import create_sql_agent
from langchain.agents.agent_toolkits import SQLDatabaseToolkit
from langchain.sql_database import SQLDatabase
from langchain import OpenAI

db = SQLDatabase.from_databricks(catalog="samples", schema="nyctaxi")
llm = OpenAI(model_name="gpt-3.5-turbo-instruct", temperature=.7)
toolkit = SQLDatabaseToolkit(db=db, llm=llm)
agent = create_sql_agent(llm=llm, toolkit=toolkit, verbose=True)

agent.run("What is the longest trip distance and how long did it take?")

次のノートブックでは、データベース内のデータをより深く理解できるように、Databricks SQL エージェントを作成および使用する方法を示します。

LangChain を使用して SQL データベース ノートブックと対話する

ノートブックを入手

Azure Databricks で提供されるモデルを LLM としてラップする

Azure Databricks で作成した LLM がある場合は、OpenAI、HuggingFace、またはその他の LLM プロバイダーの代わりに、LangChain 内でそれを直接使用することができます。

この統合では、次の 2 つのエンドポイントの種類をサポートしています。

• 実稼働および開発に推奨されるモデル提供エンドポイント。
• クラスター ドライバー プロキシ アプリ。対話型開発に推奨されます。

モデル提供エンドポイントをラップする

Azure Databricks エンドポイントを LangChain の LLM としてラップできます。 モデル提供エンドポイントを LangChain の LLM としてラップするには、次のものが必要です。

• Azure Databricks モデル提供エンドポイントにデプロイされた登録済み LLM。
• エンドポイントに対する CAN QUERY アクセス許可。

多くの場合、モデルでは、temperature や max_tokens などの重要なパラメーターが必要か、推奨されます。 次の例は、falcon-7b-instruct という名前のデプロイ済みモデルでこれらのパラメーターを入力する方法を示しています。 その他の詳細については、LangChain ドキュメント「提供エンドポイントをラップする」を参照してください。

from langchain.llms import Databricks

llm = Databricks(endpoint_name="falcon-7b-instruct", model_kwargs={"temperature": 0.1, "max_tokens": 100})
llm("How are you?")

クラスター ドライバー プロキシ アプリケーションをラップする

クラスター ドライバー プロキシ アプリケーションを LangChain の LLM としてラップするには、次のものが必要です。

• LLM は、“シングル ユーザー“ または “非分離共有“ モードで Azure Databricks 対話型クラスターに読み込まれます。
• ドライバー ノード上で実行されているローカル HTTP サーバーは、JSON 入出力を備えた HTTP POST を使用して “/” でモデルを提供します。
• アプリは [3000, 8000] の間のポート番号を使用し、ドライバーの IP アドレス、または localhost のみではなく単純に 0.0.0.0 をリッスンします。
• クラスターへの [アタッチ可能] アクセス許可。

例については、LangChain ドキュメント「クラスター ドライバー プロキシ アプリをラップする」を参照してください。