コードで外部ベクターデータベースを作成する¶
以下のノートブックでは、DataRobotのPythonクライアントを使用して外部ベクターデータベースを構築、検証、およびDataRobotプラットフォームに登録する方法について説明します。このノートブックはDataRobot Notebooksで使用するように設計されているため、このノートブックをダウンロードし、プラットフォームにアップロードして使用することを推奨します。
セットアップ¶
以下の手順では、外部ベクターデータベースをDataRobotプラットフォームと連携させるために必要な設定について説明します。
このワークフローは以下の機能フラグを使用します。これらの機能を有効にするには、DataRobotの担当者または管理者にお問い合わせください。
- Notebooksでファイルシステム管理を有効にする
- プロキシモデルを有効にする
- すべてのカスタムモデルでパブリックネットワークへのアクセスを有効にする
- 生成モデルの監視サポートを有効にする
- Enable Custom Inference Models
ノートブックサイドバーで、このノートブックのノートブックファイルシステムを有効にします。
ノートブックのセッションタイムアウトを180分に設定します。
少なくとも"Medium"(16GB RAM)インスタンスを使用して、ノートブックコンテナを再起動します。
ドキュメントアーカイブをノートブックファイルシステムにアップロードします(オプション)。
ライブラリのインストール次のライブラリをインストールします。¶
# Upgrade pip to fix langchain installation issues
!pip install --upgrade pip setuptools
!pip install "langchain" \
"langchain-community" \
"langchain-chroma" \
"sentence-transformers==3.0.0" \
"datarobotx" \
"unstructured" \
"pysqlite3-binary"
# replace sqlite3 with pysqlite3 to fix chroma issues
__import__('pysqlite3')
import sys
sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')
import datarobot as dr
import datarobotx as drx
from datarobot.models.genai.vector_database import CustomModelVectorDatabaseValidation
from datarobot.models.genai.vector_database import VectorDatabase
DataRobotに接続する PythonクライアントからDataRobotに接続するためのさまざまなオプションの詳細を参照してください。¶
Download sample dataこの例では、DataRobotの英語ドキュメントから作成されたサンプルデータセットを参照しています。¶
独自のデータを試すには、独自のローカルデータセットを参照して、このセクションおよび「テキストの読み込みと分割」セクションを変更します。
注:セルフマネージドAIプラットフォームでは、app.datarobot.com
を参照するコードサンプルは、インスタンスに適したURLに変更する必要があります。
import requests, zipfile, io
SOURCE_DOCUMENTS_ZIP_URL = "https://s3.amazonaws.com/datarobot_public_datasets/ai_accelerators/datarobot_english_documentation_5th_December.zip"
UNZIPPED_DOCS_DIR = "datarobot_english_documentation"
STORAGE_DIR = "storage"
r = requests.get(SOURCE_DOCUMENTS_ZIP_URL)
z = zipfile.ZipFile(io.BytesIO(r.content))
z.extractall(f"{STORAGE_DIR}/")
テキストの読み込みと分割¶
次に、DataRobotのドキュメントデータセットをロードし、チャンクに分割します。このレシピを別のユースケースに適用する場合は、以下の点に注意してください。
- 追加または代替のドキュメントローダーを使用します。
- 余分で不要なドキュメントを除外します。
- 適切な
chunk_size
とoverlap
を選択しします。これらはトークンではなく、文字数でカウントされます。
import re
from langchain.document_loaders import DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
SOURCE_DOCUMENTS_DIR = f"{STORAGE_DIR}/{UNZIPPED_DOCS_DIR}/"
SOURCE_DOCUMENTS_FILTER = "**/*.txt"
loader = DirectoryLoader(f"{SOURCE_DOCUMENTS_DIR}", glob=SOURCE_DOCUMENTS_FILTER)
splitter = RecursiveCharacterTextSplitter(
chunk_size=128,
chunk_overlap=0,
)
print(f"Loading {SOURCE_DOCUMENTS_DIR} directory")
data = loader.load()
print(f"Splitting {len(data)} documents")
docs = splitter.split_documents(data)
for doc in docs:
doc.metadata['source'] = re.sub(
rf'{STORAGE_DIR}/{UNZIPPED_DOCS_DIR}/datarobot_docs/en/(.+)\.md',
r'https://docs.datarobot.com/en/docs/\1.html',
doc.metadata['source']
)
doc.metadata["category"] = doc.metadata["source"].split("|")[-1].replace(".txt", "")
print(f"Created {len(docs)} documents")
Create a vector database from documents (using ChromaDB with metadata filtering)¶
Use the following cell to build a vector database from the DataRobot documentation dataset. Note that this notebook uses ChromaDB, an open source, in-memory vector store with metadata filtering support that is compatible with the DataRobot Notebooks. Additionally, this notebook uses the HuggingFace jina-embedding-t-en-v1
embeddings model (open source).
from datetime import datetime
from langchain_chroma import Chroma
from langchain_community.embeddings.sentence_transformer import (SentenceTransformerEmbeddings)
CHROMADB_DATA_PATH = f"{STORAGE_DIR}/chromadb"
CHROMADB_EMBEDDING_CACHE_FOLDER = STORAGE_DIR + '/sentencetransformers'
CHROMADB_EMBEDDING_FUNCTION = SentenceTransformerEmbeddings(model_name="jinaai/jina-embedding-t-en-v1", cache_folder=CHROMADB_EMBEDDING_CACHE_FOLDER)
def create_chromadb_from_documents(docs, embedding_function, persist_directory):
start_time = datetime.now()
print(f'>>> BEGIN ({start_time.strftime("%H:%M:%S")}): Creating ChromaDB from documents')
print(f'Embedding function: {embedding_function}')
print(f'ChromaDB data directory: {persist_directory}')
print(' ')
print(f'Documents for loading: {len(docs)}')
db = Chroma.from_documents(docs, embedding_function, persist_directory=persist_directory)
end_time = datetime.now()
print(' ')
print(f'>>> END ({end_time.strftime("%H:%M:%S")}): Creating ChromaDB from documents')
print(f'Loaded {len(docs)} documents.')
print(f"Chroma VectorDB now has {db._collection.count()} documents")
total_elapsed_min = (end_time - start_time).total_seconds() / 60
document_average_sec = (end_time - start_time).total_seconds() / len(docs)
print("Total Elapsed", "%.2f" % total_elapsed_min, "minutes")
print("Document Average", "%.2f" % document_average_sec, "seconds")
return db
print(f"Created {len(docs)} documents")
db = create_chromadb_from_documents(docs, CHROMADB_EMBEDDING_FUNCTION, CHROMADB_DATA_PATH)
print(db._collection.count())
ベクターデータベースのテスト¶
Use the following cell to test the vector database by having the model perform a similarity search with metadata filtering; it will return the top five documents matching the query provided.
question = "What is MLOps?"
top_k = 5
metadata_filter = {"category": {"$eq": "index"}}
results_with_scores = db.similarity_search_with_score(
question,
k=top_k,
filter=metadata_filter,
)
print(len(results_with_scores))
for doc, score in results_with_scores:
print("********************************************************************************")
print(" ")
print("----------")
print(f"METADATA: {doc.metadata}, Score: {score}")
print(" ")
print("----------")
print(f"CONTENT: {doc.page_content}")
print(" ")
非構造化カスタムモデルをデプロイするためのフックを定義する¶
以下のセルは、非構造化カスタムモデルのデプロイに使用されるメソッドを定義します。これらには、カスタムモデルのロードと、スコアリングのためのモデルの使用が含まれます。
import os
def load_model(input_dir):
"""Custom model hook for loading our knowledge base."""
import os
print("Loading model")
STORAGE_DIR = "storage"
CHROMADB_DATA_PATH = f"{STORAGE_DIR}/chromadb"
CHROMADB_EMBEDDING_CACHE_FOLDER = CHROMADB_DATA_PATH + '/sentencetransformers'
CHROMADB_EMBEDDING_MODEL_NAME = 'jinaai/jina-embedding-t-en-v1'
# https://docs.trychroma.com/troubleshooting#sqlite
__import__('pysqlite3')
import sys
sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')
from langchain_chroma import Chroma
from langchain.embeddings.sentence_transformer import SentenceTransformerEmbeddings
print(f'CHROMADB_DATA_PATH = {CHROMADB_DATA_PATH}')
print(f'CHROMADB_EMBEDDING_CACHE_FOLDER = {CHROMADB_EMBEDDING_CACHE_FOLDER}')
print(f'CHROMADB_EMBEDDING_MODEL_NAME = {CHROMADB_EMBEDDING_MODEL_NAME}')
CHROMADB_EMBEDDING_FUNCTION = SentenceTransformerEmbeddings(model_name="jinaai/jina-embedding-t-en-v1", cache_folder=CHROMADB_EMBEDDING_CACHE_FOLDER)
db = Chroma(persist_directory=CHROMADB_DATA_PATH, embedding_function=CHROMADB_EMBEDDING_FUNCTION)
print(f'Loaded ChromaDB with {db._collection.count} chunks')
return db
def score_unstructured(model, data, query, **kwargs) -> str:
"""Custom model hook for retrieving relevant docs with our knowledge base.
When requesting predictions from the deployment, pass a dictionary
with the following keys:
- 'question' the question to be passed to the vector store retriever
- 'metadata' the metadata filter to be passed to the vector store retriever
- 'top_k' the number of results to return
datarobot-user-models (DRUM) handles loading the model and calling
this function with the appropriate parameters.
Returns:
--------
rv : str
Json dictionary with keys:
- 'question' user's original question
- 'relevant' the generated answer to the question
- 'metadata' - metadata for each document
- 'error' - error message if exception in handling request
"""
import json
try:
data_dict = json.loads(data)
question = data_dict['question']
top_k = data_dict.get("k", 10)
metadata_filter = data_dict.get("filter", None)
results_with_scores = model.similarity_search_with_score(
question,
k=top_k,
filter=metadata_filter,
)
print(f'Returned {len(results_with_scores)} results')
relevant, metadata = [], []
for doc, score in results_with_scores:
relevant.append(doc.page_content)
doc.metadata["similarity_score"] = score
metadata.append(doc.metadata)
rv = {
"question": question,
"relevant": relevant,
"metadata": metadata,
}
except Exception as e:
rv = {'error': f"{e.__class__.__name__}: {str(e)}"}
return json.dumps(rv), {"mimetype": "application/json", "charset": "utf8"}
ローカルでフックをテストする¶
デプロイに進む前に、以下のセルを使用して、カスタムモデルフックが正しく機能することを確認します。
import json
# Test the hooks locally
score_unstructured(
load_model(f"{STORAGE_DIR}"),
json.dumps(
{
"question": "How do I replace a custom model on an existing custom environment?",
"filter": {"category": {"$eq": "index"}},
"k": 1,
}
),
None,
)
ナレッジベースのデプロイ¶
以下のセルでは、次のことを行う便利なメソッドを使用しています。
storage/deploy/
の内容を含む新しいカスタムモデル環境を構築する。- 提供されたフックで新しいカスタムモデルを構築する。
- 非構造化カスタムモデルをDataRobotにデプロイする。
- 予測に使用できるオブジェクトを返す。
この例では、事前構築済みの環境を使用します。
また、environment_id
を指定して、既存のカスタムモデル環境を代わりに使用することで、カスタムモデルフックの反復サイクルを短くすることもできます。
DataRobotのカスタムモデルワークショップからアカウントの既存の構築済み環境を参照してください。
dr.ExecutionEnvironment.list("Python 3.11")
[ExecutionEnvironment('[DataRobot][NVIDIA] Python 3.11 GenAI'), ExecutionEnvironment('[DataRobot] Python 3.11 GenAI'), ExecutionEnvironment('[GenAI] Python 3.11 with Moderations')]
deployment = drx.deploy(
model="storage/chromadb/",
name="External DR Knowledge Base",
hooks={
"score_unstructured": score_unstructured,
"load_model": load_model
},
extra_requirements=["langchain_chroma", "pysqlite3-binary"],
# Re-use existing environment if you want to change the hook code,
# and not requirements
environment_id=dr.ExecutionEnvironment.list("Python 3.11")[-1].id,
)
デプロイのテスト¶
デプロイが質問に対して正常に回答できるかどうかをテストします。
#deployment = drx.Deployment("ADD_VALUE_HERE")
deployment.predict_unstructured(
{
"question": "How do I replace a custom model on an existing custom environment?",
"filter": {"category": {"$eq": "index"}},
"k": 1,
}
)
外部ベクターデータベースからのベクターデータベースの検証および作成これらのメソッドは外部ベクターデータベースを実行、検証、統合します。¶
この例では、ユースケースと検証を関連付け、そのユースケース内にベクターデータベースを作成します。
use_case_id
を設定して既存のユースケースを指定するか、その名前で新しいユースケースを作成します。
use_case_id = "ADD_VALUE_HERE"
use_case = dr.UseCase.get(use_case_id)
# UNCOMMENT if you wish to create a new UseCase
# use_case = dr.UseCase.create()
CustomModelVectorDatabaseValidation.create
は、ベクターデータベースの検証を実行します。デプロイIDを必ず入力します。
external_vdb_validation = CustomModelVectorDatabaseValidation.create(
prompt_column_name="question",
target_column_name="relevant",
deployment_id=deployment.dr_deployment.id,
use_case=use_case,
wait_for_completion=True
)
assert external_vdb_validation.validation_status == "PASSED"
検証が完了したら、VectorDatabase.create_from_custom_model()
を使用してベクターデータベースを連携します。ユースケース名(またはユースケースID)、外部ベクターデータベースの名前、および前のセルから返される検証IDを指定する必要があります。
vdb = VectorDatabase.create_from_custom_model(
name="DR External Vector Database",
use_case=use_case,
validation_id=external_vdb_validation.id
)
assert vdb.execution_status == "COMPLETED"
print(f"Vector Database ID: {vdb.id}")
これで、このベクターデータベースIDを GenAI E2E基本ステップで使用して、外部ベクターデータベースでLLMブループリントを作成できるようになりました。