LinkedInでこちらの記事を見かけました。

LlamaIndexのRAGのLLMにDBRXを指定できるとな！

やってみます。ベースにしているのはこちらの記事のコードです。

マニュアルはこちらですが、所々コードを修正しています。

注意
この記事で説明しているコードを動作させるには、対象のDatabricksワークスペースでDatabricks 基盤モデル APIが利用できるようになっている必要があります。

ライブラリのインストール

llama-index-llms-databricksをインストールします。

%pip install -U llama-index llama-index-llms-databricks llama-index-core
%pip install "databricks-sql-connector[sqlalchemy]"
dbutils.library.restartPython()

Unity Catalogとの接続

from llama_index.core import download_loader, SQLDatabase, VectorStoreIndex, SimpleDirectoryReader, Settings

DatabaseReader = download_loader('DatabaseReader')

access_token    = dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiToken().getOrElse(None)
server_hostname = "<Databricksワークスペースのホスト名>"
http_path       = "sql/protocolv1/o/2556758628403379/1118-004519-nu0u899x"
catalog         = "takaakiyayoi_catalog"
schema          = "qiita_2023"

reader = DatabaseReader(
    sql_database = SQLDatabase.from_uri(f"databricks://token:{access_token}@{server_hostname}?" +
        f"http_path={http_path}&catalog={catalog}&schema={schema}")
)

テキストデータのロード

query = f"""
SELECT
    body
FROM takaakiyayoi_catalog.qiita_2023.taka_qiita_2023
"""

documents = reader.load_data(query=query)

インデックスの作成

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

index = VectorStoreIndex.from_documents(documents)

LLMの設定

DBRXを使うようにします。

from llama_index.llms.databricks import Databricks

dbrx = Databricks(
    model="databricks-dbrx-instruct",
    api_key="<Databricksのパーソナルアクセストークン>",
    api_base="https://<Databricksワークスペースのホスト名>/serving-endpoints/",
)

# LLMを定義
Settings.llm = dbrx

RAGへの問い合わせ

query_engine = index.as_query_engine()
response = query_engine.query("Sparkとは")
print("回答: ", response)

サクッと返ってきましたが英語です。

回答: Spark is a powerful data processing engine that is significantly faster than Hadoop MapReduce due to its ability to store intermediate results in memory. It offers APIs for various libraries, including machine learning (MLlib), interactive queries (Spark SQL), real-time data processing (Structured Streaming), and graph processing (GraphX). Spark's design philosophy focuses on speed, ease of use, modularity, and extensibility. It provides a simple programming model based on transformations and actions, supports multiple programming languages, and integrates with various storage systems and data sources. Spark's DataFrameReader and DataFrameWriter can be extended to load data from various sources like Apache Kafka, Kinesis, Azure Storage, and Amazon S3. By using a managed and governed platform like Databricks Lakehouse, developers can focus on building products and improving their productivity instead of managing Spark infrastructure.

プロンプトをカスタマイズします。

from llama_index.core import PromptTemplate

template = (
    "あなたは有能なアシスタントです。以下のコンテキストがあります。\n"
    "---------------------\n"
    "{context_str}"
    "\n---------------------\n"
    "この情報を用いて質問に日本語で回答してください: {query_str}\n"
)
qa_template = PromptTemplate(template)

query_engine = index.as_query_engine(text_qa_template=qa_template)
response = query_engine.query("Sparkとは")
print("回答: ", response)

回答: Sparkは、大規模データ処理のための高速で汎用的なクラスターリソース管理システムです。Sparkは、Hadoop MapReduceよりもはるかに高速で、機械学習、インタラクティブなクエリー、リアルタイムデータ処理、グラフ処理など、さまざまなワークロードに対応するAPIを持つライブラリと連携することができます。Sparkは、中間結果をメモリに保持することで高速化を実現し、ユーザーにシンプルなプログラミングモデルを提供することで使いやすさを実現しています。Sparkは、Scala、Java、SQL、Rなどのプログラミング言語で表現することができ、さまざまなストレージシステムと統合することができます。

response = query_engine.query("PrivateLink構成を組むにはどうしたらいいですか")
print("回答: ", response)

回答: PrivateLink構成を組むには、以下の手順を実行します。

Azure Databricks または AWS Databricks のワークスペースを作成します。

Azure Private Link または AWS PrivateLink を有効にし、Databricks ワークスペースと接続します。

Private Link エンドポイントを作成し、Databricks ユーザーとコントロールプレーン、コントロールプレーンとデータプレーンの間に接続します。

プライベートネットワークでトラフィックをルーティングするようにネットワークを構成します。

必要に応じて、コンプライアンス要件を満たすために、データ周辺にセキュアな防御線を構築します。

これにより、エンドツーエンドのプライベートネットワークを確立し、データ漏洩のリスクを軽減することができます。

LlamaIndexでDBRXを使う