コードで外部ベクターデータベースを作成する¶

!pip install "langchain==0.0.244" \
             "faiss-cpu==1.7.4" \
             "sentence-transformers==2.2.2" \
             "datarobotx==0.1.25"

import datarobot as dr
import datarobotx as drx
from datarobot.models.genai.vector_database import CustomModelVectorDatabaseValidation
from datarobot.models.genai.vector_database import VectorDatabase

endpoint = "https://app.datarobot.com/api/v2"
token="<ADD_VALUE_HERE>"
dr.Client(endpoint=endpoint, token=token)
drx.Context(token=token, endpoint=endpoint)

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}/")

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=2000,
    chunk_overlap=1000,
)

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']
    )
print(f"Created {len(docs)} documents")

from langchain.embeddings.sentence_transformer import SentenceTransformerEmbeddings
from langchain.vectorstores.faiss import FAISS
from langchain.docstore.document import Document
import torch

if not torch.cuda.is_available():
    EMBEDDING_MODEL_NAME = "all-MiniLM-L6-v2"
else:
    EMBEDDING_MODEL_NAME = "all-mpnet-base-v2"

# Will download the model the first time it runs
embedding_function = SentenceTransformerEmbeddings(
    model_name=EMBEDDING_MODEL_NAME,
    cache_folder="storage/deploy/sentencetransformers",
)
try:
    # Load existing db from disk if previously built
    db = FAISS.load_local("storage/deploy/faiss-db", embedding_function)
except:
    texts = [doc.page_content for doc in docs]
    metadatas = [doc.metadata for doc in docs]   
    # Build and save the FAISS db to persistent notebook storage; this can take some time w/o GPUs
    db = FAISS.from_texts(texts, embedding_function, metadatas=metadatas)  
    db.save_local("storage/deploy/faiss-db")

print(f"FAISS VectorDB has {db.index.ntotal} documents")

db.similarity_search("How do I replace a custom model on an existing custom environment?")
#db.max_marginal_relevance_search("How do I replace a custom model on an existing custom environment?")

import os


def load_model(input_dir):
    """Custom model hook for loading our knowledge base."""
    import os
    from langchain.embeddings.sentence_transformer import SentenceTransformerEmbeddings
    from langchain.vectorstores.faiss import FAISS
    embedding_function = SentenceTransformerEmbeddings(
        model_name=EMBEDDING_MODEL_NAME,
        cache_folder=input_dir + '/' + 'storage/deploy/sentencetransformers',
    )
    db = FAISS.load_local(input_dir + "/" + "storage/deploy/faiss-db", embedding_function)
    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

    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
            - 'references' list of references that were used to generate the answer
            - 'error' - error message if exception in handling request
    """
    import json
    from langchain.vectorstores.base import VectorStoreRetriever
    try:
        db = model
        data_dict = json.loads(data)
        retriever = VectorStoreRetriever(vectorstore=db)
        documents = retriever.get_relevant_documents(data_dict['question'])
        relevant_text_list = [doc.page_content for doc in documents]
        references = [doc.metadata['source'] for doc in documents]
        rv = {
            "question": data_dict["question"],
            "relevant": relevant_text_list,
            "references": references
        }
    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("."),
    json.dumps(
        {
            "question": "How do I replace a custom model on an existing custom environment?",
        }
	),
    None,
)

deployment = drx.deploy(
    model="storage/deploy/",
    name="External DR Knowledge Base",
    hooks={
        "score_unstructured": score_unstructured,
        "load_model": load_model
    },
    # extra_requirements=["langchain", "faiss-cpu", "sentence-transformers", "openai"],
    # Re-use existing environment if you want to change the hook code,
    # and not requirements
    environment_id=dr.ExecutionEnvironment.list("Python 3.9 GenAI")[0].id,
)

deployment.predict_unstructured(
    {
        "question": "How do I replace a custom model on an existing custom environment?",
    }
)

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()

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"

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}")