生成AIに関するドキュメント > コードベースのGenAI > BYO埋め込みからのベクターデータベースの作成

BYO埋め込みからのベクターデータベースの作成¶

以下のノートブックでは、BYO（bring-your-own）埋め込みからベクターデータベースを構築して検証し、Pythonクライアントを使用して、そのベクターデータベースを登録する方法について説明します。このノートブックは、DataRobot Notebooksで使用するように設計されているため、ダウンロードし、プラットフォームにアップロードして使用することを推奨します。

セットアップ¶

以下の手順では、ベクターデータベースをDataRobotプラットフォームと連携させるために必要な設定について説明します。

このワークフローは以下の機能フラグを使用します。これらの機能を有効にするには、DataRobotの担当者または管理者にお問い合わせください。

Notebooksでファイルシステム管理を有効にする
すべてのカスタムモデルでパブリックネットワークへのアクセスを有効にする
生成モデルの監視サポートを有効にする
Enable Custom Inference Models
Enable GenAI Experimentation

ノートブックサイドバーで、このノートブックのノートブックファイルシステムを有効にします。
ノートブックのセッションタイムアウトを180分に設定します。
少なくとも"Medium"（16GB RAM）インスタンスを使用して、ノートブックコンテナを再起動します。
ドキュメントアーカイブをノートブックファイルシステムにアップロードします（オプション）。

要件をインストールする¶

DataRobotプラットフォーム、予測環境、およびベクターデータベース機能とのインターフェイスとして、以下のライブラリとモジュールをインポートします。

In [55]:

Copied!





import datarobot as dr
from datarobot.enums import PredictionEnvironmentPlatform
from datarobot.enums import PredictionEnvironmentModelFormats
from datarobot.models.genai.custom_model_embedding_validation import CustomModelEmbeddingValidation
from datarobot.models.genai.vector_database import VectorDatabase
from datarobot.models.genai.vector_database import ChunkingParameters
import datarobot as dr
from datarobot.enums import PredictionEnvironmentPlatform
from datarobot.enums import PredictionEnvironmentModelFormats
from datarobot.models.genai.custom_model_embedding_validation import CustomModelEmbeddingValidation
from datarobot.models.genai.vector_database import VectorDatabase
from datarobot.models.genai.vector_database import ChunkingParameters

DataRobotに接続する¶

DataRobot PythonクライアントからDataRobotに接続するために、DataRobotのendpointとtokenを指定します。PythonクライアントからDataRobotに接続するためのオプションの詳細をご確認ください。

In [ ]:

Copied!

endpoint="https://app.datarobot.com/api/v2"
token="ADD_TOKEN_HERE>"
dr.Client(endpoint=endpoint, token=token)
endpoint="https://app.datarobot.com/api/v2"
token="ADD_TOKEN_HERE>"
dr.Client(endpoint=endpoint, token=token)

Out[ ]:

<datarobot.rest.RESTClientObject at 0x17e7f2d50>

Select an environment¶

dr.ExecutionEnvironment.list()を使って、組織で利用可能な環境を順に確認し、ベクターデータベースの基本環境として[DataRobot] Python 3.11 GenAIという名前の環境を選択して、base_environmentに割り当てます。

In [57]:

Copied!





execution_environments = dr.ExecutionEnvironment.list()

base_environment = None
environment_versions = None

for execution_environment in execution_environments:
    # print(execution_environment)
    if execution_environment.name == "[DataRobot] Python 3.11 GenAI":
        base_environment = execution_environment
        environment_versions = dr.ExecutionEnvironmentVersion.list(
            execution_environment.id
        )
        break

environment_version = environment_versions[0]
print(base_environment)
print(environment_version)
execution_environments = dr.ExecutionEnvironment.list()

base_environment = None
environment_versions = None

for execution_environment in execution_environments:
    # print(execution_environment)
    if execution_environment.name == "[DataRobot] Python 3.11 GenAI":
        base_environment = execution_environment
        environment_versions = dr.ExecutionEnvironmentVersion.list(
            execution_environment.id
        )
        break

environment_version = environment_versions[0]
print(base_environment)
print(environment_version)

ExecutionEnvironment('[DataRobot] Python 3.11 GenAI')
ExecutionEnvironmentVersion('v23')

カスタム埋め込みモデルを作成¶

dr.CustomInferenceModel.list()を使って、利用可能なカスタムモデルからall-MiniLM-L6-v2-embedding-modelを検索します。このカスタムモデルが存在しない場合は、dr.CustomInferenceModel.create()を使ってcustom_modelとして作成します。このカスタムモデルが存在する場合は、custom_modelに割り当てます。

In [58]:

Copied!





CUSTOM_MODEL_NAME = "all-MiniLM-L6-v2-embedding-model"
if CUSTOM_MODEL_NAME not in [c.name for c in dr.CustomInferenceModel.list()]:
    # Create a new custom model
    print("Creating new custom model")
    custom_model = dr.CustomInferenceModel.create(
        name=CUSTOM_MODEL_NAME,
        target_type=dr.TARGET_TYPE.UNSTRUCTURED,
        is_training_data_for_versions_permanently_enabled=True
    )
else:
    print("Custom Model Exists")
    custom_model = [c for c in dr.CustomInferenceModel.list() if c.name == CUSTOM_MODEL_NAME].pop()
CUSTOM_MODEL_NAME = "all-MiniLM-L6-v2-embedding-model"
if CUSTOM_MODEL_NAME not in [c.name for c in dr.CustomInferenceModel.list()]:
    # Create a new custom model
    print("Creating new custom model")
    custom_model = dr.CustomInferenceModel.create(
        name=CUSTOM_MODEL_NAME,
        target_type=dr.TARGET_TYPE.UNSTRUCTURED,
        is_training_data_for_versions_permanently_enabled=True
    )
else:
    print("Custom Model Exists")
    custom_model = [c for c in dr.CustomInferenceModel.list() if c.name == CUSTOM_MODEL_NAME].pop()

Creating new custom model

カスタムコードのディレクトリを作成¶

custom_embedding_modelというディレクトリを作成し、カスタム埋め込みモデルのコードを記述します。

In [ ]:

Copied!

import os
os.mkdir('custom_embedding_model')
import os
os.mkdir('custom_embedding_model')

カスタム埋め込みモデルのコードを書く¶

カスタム埋め込みモデルのコードをcustom.pyファイルに記述し、all-MiniLM-L6-v2から非構造化モデルを作成します。

In [ ]:

Copied!





%%writefile ./custom_embedding_model/custom.py
from sentence_transformers import SentenceTransformer


def load_model(input_dir):
    return SentenceTransformer("all-MiniLM-L6-v2")


def score_unstructured(model, data, query, **kwargs):
    import json

    data_dict = json.loads(data)
    outputs = model.encode(data_dict["input"])
    return json.dumps(
        {
            "result": outputs.tolist(),
            "device": str(model._target_device)
        }
    )
%%writefile ./custom_embedding_model/custom.py
from sentence_transformers import SentenceTransformer


def load_model(input_dir):
    return SentenceTransformer("all-MiniLM-L6-v2")


def score_unstructured(model, data, query, **kwargs):
    import json

    data_dict = json.loads(data)
    outputs = model.encode(data_dict["input"])
    return json.dumps(
        {
            "result": outputs.tolist(),
            "device": str(model._target_device)
        }
    )

Writing ./custom_embedding_model/custom.py

カスタムモデルのバージョンを作成¶

dr.CustomModelVersion.create_cleanを使い、前の手順で定義したcustom_model、base_environment、filesでカスタムモデルバージョンを作成します。さらに、dr.NETWORK_EGRESS_POLICY.PUBLICを使って、パブリックネットワークへのアクセスを有効にします。

In [ ]:

Copied!





# Create a new custom model version in DataRobot
print("Upload new version of model to DataRobot")
model_version = dr.CustomModelVersion.create_clean(
    custom_model_id=custom_model.id,
    base_environment_id=base_environment.id,
    files=[("./custom_embedding_model/custom.py", "custom.py")],
    network_egress_policy=dr.NETWORK_EGRESS_POLICY.PUBLIC,
)
# Create a new custom model version in DataRobot
print("Upload new version of model to DataRobot")
model_version = dr.CustomModelVersion.create_clean(
    custom_model_id=custom_model.id,
    base_environment_id=base_environment.id,
    files=[("./custom_embedding_model/custom.py", "custom.py")],
    network_egress_policy=dr.NETWORK_EGRESS_POLICY.PUBLIC,
)

Upload new version of model to DataRobot

カスタムモデルの登録¶

dr.RegisteredModel.list()を使って、利用可能なカスタムモデルから（CUSTOM_MODEL_NAMEに割り当てられた）all-MiniLM-L6-v2-embedding-modelを検索します。登録されたモデルが存在しない場合は、dr.RegisteredModelVersion.create_for_custom_model_versionを使用してregistered_model_versionを作成します。登録されたモデルが存在する場合は、registered_modelに割り当て、dr.RegisteredModelVersion.create_for_custom_model_versionを使用してregistered_model_versionを作成します。

In [ ]:

Copied!





if CUSTOM_MODEL_NAME not in [m.name for m in dr.RegisteredModel.list()]:
    print("Creating New Registered Model")
    registered_model_version = dr.RegisteredModelVersion.create_for_custom_model_version(
        model_version.id,
        name=CUSTOM_MODEL_NAME,
        registered_model_name=CUSTOM_MODEL_NAME
    )
else:
    print("Using Existing Model")
    registered_model = [m for m in dr.RegisteredModel.list() if m.name == CUSTOM_MODEL_NAME].pop()
    registered_model_version = dr.RegisteredModelVersion.create_for_custom_model_version(
        model_version.id,
        name=CUSTOM_MODEL_NAME,
        registered_model_id=registered_model.id
    )
if CUSTOM_MODEL_NAME not in [m.name for m in dr.RegisteredModel.list()]:
    print("Creating New Registered Model")
    registered_model_version = dr.RegisteredModelVersion.create_for_custom_model_version(
        model_version.id,
        name=CUSTOM_MODEL_NAME,
        registered_model_name=CUSTOM_MODEL_NAME
    )
else:
    print("Using Existing Model")
    registered_model = [m for m in dr.RegisteredModel.list() if m.name == CUSTOM_MODEL_NAME].pop()
    registered_model_version = dr.RegisteredModelVersion.create_for_custom_model_version(
        model_version.id,
        name=CUSTOM_MODEL_NAME,
        registered_model_id=registered_model.id
    )

Creating New Registered Model

埋め込みモデル用の予測環境を作成します。¶

dr.PredictionEnvironment.list()を使って、利用可能な予測環境からBYO埋め込みモデルの予測環境を検索します。その予測環境が存在しない場合は、dr.PredictionEnvironment.createを使用してDataRobotサーバーレスprediction_environmentを作成します。その予測環境が存在する場合は、prediction_environmentに割り当てます。

In [ ]:

Copied!





PREDICTION_ENVIRONMENT_NAME = "Prediction environment for BYO embeddings models"

prediction_environment = None
for _prediction_environment in dr.PredictionEnvironment.list():
    if _prediction_environment.name == PREDICTION_ENVIRONMENT_NAME:
        prediction_environment = _prediction_environment

if prediction_environment is None:
    prediction_environment = dr.PredictionEnvironment.create(
        name=PREDICTION_ENVIRONMENT_NAME,
        platform=PredictionEnvironmentPlatform.DATAROBOT_SERVERLESS,
        supported_model_formats=[
            PredictionEnvironmentModelFormats.DATAROBOT,
            PredictionEnvironmentModelFormats.CUSTOM_MODEL
        ],
    )
PREDICTION_ENVIRONMENT_NAME = "Prediction environment for BYO embeddings models"

prediction_environment = None
for _prediction_environment in dr.PredictionEnvironment.list():
    if _prediction_environment.name == PREDICTION_ENVIRONMENT_NAME:
        prediction_environment = _prediction_environment

if prediction_environment is None:
    prediction_environment = dr.PredictionEnvironment.create(
        name=PREDICTION_ENVIRONMENT_NAME,
        platform=PredictionEnvironmentPlatform.DATAROBOT_SERVERLESS,
        supported_model_formats=[
            PredictionEnvironmentModelFormats.DATAROBOT,
            PredictionEnvironmentModelFormats.CUSTOM_MODEL
        ],
    )

カスタム埋め込みモデルをデプロイします。¶

dr.Deployment.list()を使って、利用可能なデプロイからall-MiniLM-L6-v2のデプロイを検索します。予測環境が存在しない場合は、deploymentを作成し、前の手順においてdr.Deployment.create_from_registered_model_versionで作成したregistered_model_versionをデプロイします。デプロイが存在する場合は、deploymentに割り当てます。

In [ ]:

Copied!





MODEL_DEPLOYMENT_NAME = "Deployment for all-MiniLM-L6-v2"

if MODEL_DEPLOYMENT_NAME not in [d.label for d in dr.Deployment.list()]:
    deployment = dr.Deployment.create_from_registered_model_version(
        registered_model_version.id,
        label=MODEL_DEPLOYMENT_NAME,
        max_wait=1000,
        prediction_environment_id=prediction_environment.id
    )
else:
    deployment = [d for d in dr.Deployment.list() if d.label == MODEL_DEPLOYMENT_NAME]
MODEL_DEPLOYMENT_NAME = "Deployment for all-MiniLM-L6-v2"

if MODEL_DEPLOYMENT_NAME not in [d.label for d in dr.Deployment.list()]:
    deployment = dr.Deployment.create_from_registered_model_version(
        registered_model_version.id,
        label=MODEL_DEPLOYMENT_NAME,
        max_wait=1000,
        prediction_environment_id=prediction_environment.id
    )
else:
    deployment = [d for d in dr.Deployment.list() if d.label == MODEL_DEPLOYMENT_NAME]

BYO埋め込みのユースケースを作成する¶

dr.UseCase.createを使って、ベクターデータベースを使用するユースケースを作成し、use_caseに割り当てます。

In [ ]:

Copied!

use_case = dr.UseCase.create(name="For BYO embeddings")
use_case = dr.UseCase.create(name="For BYO embeddings")

ユースケースにデータセットをアップロードする¶

dr.Dataset.create_from_urlを使って、ベクターデータベースのサンプルデータセットをアップロードし、datasetに割り当てます。

In [ ]:

Copied!





# this can be updated with any public URL that is pointing to a .zip file
# in the expected format
dataset_url = "https://s3.amazonaws.com/datarobot_public_datasets/genai/pirate_resumes.zip"

# We will use a vector database with our GenAI models. Let's upload a dataset with our documents.
# If you wish to use a local file as dataset, change this to
# `dataset = dr.Dataset.create_from_file(local_file_path)`
dataset = dr.Dataset.create_from_url(dataset_url)
# this can be updated with any public URL that is pointing to a .zip file
# in the expected format
dataset_url = "https://s3.amazonaws.com/datarobot_public_datasets/genai/pirate_resumes.zip"

# We will use a vector database with our GenAI models. Let's upload a dataset with our documents.
# If you wish to use a local file as dataset, change this to
# `dataset = dr.Dataset.create_from_file(local_file_path)`
dataset = dr.Dataset.create_from_url(dataset_url)

Then, add the dataset to the use_case created in the previous section.

In [ ]:

Copied!

# Attach dataset to use case.
use_case.add(dataset)
# Attach dataset to use case.
use_case.add(dataset)

Out[ ]:

<datarobot.models.use_cases.use_case.UseCaseReferenceEntity at 0x17e4b6150>

カスタム埋め込みモデルの検証と作成¶

CustomModelVectorDatabaseValidation.create関数は、ベクターデータベースの検証を実行して、必要な設定を行い、このノートブックの前の手順で作成したuse_caseおよびdeploymentにカスタム埋め込みモデルを関連付けます。この手順では、検証IDがcustom_model_embedding_validationに保存されます。

In [ ]:

Copied!





# Create BYO embeddings validation using prepared deployment
custom_model_embedding_validation = CustomModelEmbeddingValidation.create(
    prompt_column_name="input",
    target_column_name="result",
    deployment_id=deployment.id,
    use_case = use_case,
    name="BYO embeddings",
    wait_for_completion=True,
    prediction_timeout=300,
)
# Create BYO embeddings validation using prepared deployment
custom_model_embedding_validation = CustomModelEmbeddingValidation.create(
    prompt_column_name="input",
    target_column_name="result",
    deployment_id=deployment.id,
    use_case = use_case,
    name="BYO embeddings",
    wait_for_completion=True,
    prediction_timeout=300,
)

チャンキングパラメーターを設定して、ベクターデータベースを作成します¶

検証が完了したら、ChunkingParameter()を設定し、VectorDatabase.create()を使用してベクターデータベースを統合します。この手順では、前のセクションで定義したcustom_model_embedding_validation、datasetおよびuse_caseを使用します。

In [ ]:

Copied!





# Use created validation to set up chunking parameters
chunking_parameters = ChunkingParameters(
    embedding_validation=custom_model_embedding_validation,
    chunking_method="recursive",
    chunk_size=256,
    chunk_overlap_percentage=50,
    separators=["\n\n", "\n", " ", ""],
    embedding_model=None,
)


vdb = VectorDatabase.create(
    dataset_id=dataset.id,
    chunking_parameters=chunking_parameters,
    use_case=use_case
)
# Use created validation to set up chunking parameters
chunking_parameters = ChunkingParameters(
    embedding_validation=custom_model_embedding_validation,
    chunking_method="recursive",
    chunk_size=256,
    chunk_overlap_percentage=50,
    separators=["\n\n", "\n", " ", ""],
    embedding_model=None,
)


vdb = VectorDatabase.create(
    dataset_id=dataset.id,
    chunking_parameters=chunking_parameters,
    use_case=use_case
)

In [ ]:

Copied!

vdb = VectorDatabase.get(vdb.id)
assert vdb.execution_status == "COMPLETED"
vdb = VectorDatabase.get(vdb.id)
assert vdb.execution_status == "COMPLETED"