ChromaDBベクターデータベースの構築とホスト¶

# Upgrade pip to fix langchain installation issues
!pip install --upgrade pip "setuptools<82"

!pip install "langchain" \
             "langchain-community" \
             "langchain-chroma" \
             "langchain-text-splitters" \
             "sentence-transformers==3.0.0" \
             "datarobot" \
             "datarobot-predict" \
             "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
from datarobot.models.genai.vector_database import CustomModelVectorDatabaseValidation
from datarobot.models.genai.vector_database import VectorDatabase

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_community.document_loaders import DirectoryLoader, TextLoader
from langchain_text_splitters 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,
    loader_cls=TextLoader,
)
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")

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

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

def load_model(input_dir):

    """Custom model hook for loading our knowledge base."""
    import os
    print("Loading model")
    
    chromadb_data_path = os.path.join(input_dir, "chromadb")
    chromadb_embedding_cache_folder = os.path.join(input_dir, "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_community.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=chromadb_embedding_model_name,
        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, **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,
        }
    ),
)

import inspect
from pathlib import Path

CUSTOM_MODEL_NAME = "External DR Knowledge Base using ChromaDB"
DEPLOY_DIR = Path("storage/deploy_model")
DEPLOY_DIR.mkdir(parents=True, exist_ok=True)

(DEPLOY_DIR / "custom.py").write_text(
    inspect.getsource(load_model) + "\n\n" + inspect.getsource(score_unstructured)
)
(DEPLOY_DIR / "requirements.txt").write_text(
    "langchain_chroma\n"
    "pysqlite3-binary\n"
    "langchain-community\n"
    "sentence-transformers==3.0.0\n"
)

model_files = [
    (str(DEPLOY_DIR / "custom.py"), "custom.py"),
    (str(DEPLOY_DIR / "requirements.txt"), "requirements.txt"),
]
for artifact_dir in ("chromadb", "sentencetransformers"):
    source_root = Path("storage") / artifact_dir
    for path in source_root.rglob("*"):
        if path.is_file():
            model_files.append(
                (str(path), f"{artifact_dir}/{path.relative_to(source_root).as_posix()}")
            )

genai_environment = dr.ExecutionEnvironment.list(
    search_for="[GenAI] Python 3.12 with Moderations"
)[0]

existing_models = [m for m in dr.CustomInferenceModel.list() if m.name == CUSTOM_MODEL_NAME]
if existing_models:
    custom_model = existing_models[0]
else:
    custom_model = dr.CustomInferenceModel.create(
        name=CUSTOM_MODEL_NAME,
        target_type=dr.TARGET_TYPE.UNSTRUCTURED,
        language="python",
    )

print("Uploading custom model version")
model_version = dr.CustomModelVersion.create_clean(
    custom_model_id=custom_model.id,
    base_environment_id=genai_environment.id,
    files=model_files,
    maximum_memory=8 * 1024 * 1024 * 1024,
    max_wait=60 * 30,
)

print("Building custom model dependency image")
try:
    build_info = dr.CustomModelVersionDependencyBuild.start_build(
        custom_model_id=custom_model.id,
        custom_model_version_id=model_version.id,
        max_wait=60 * 60,
    )
except dr.errors.ClientError as e:
    if "already has a dependency image" in str(e):
        print("Dependency build already exists, skipping build step")
        build_info = dr.CustomModelVersionDependencyBuild.get_build_info(
            custom_model.id,
            model_version.id,
        )
    else:
        raise

if build_info is not None and build_info.build_status != "success":
    raise RuntimeError(
        f"Dependency build failed with status '{build_info.build_status}'.\n"
        f"{build_info.get_log()}"
    )

print("Creating deployment")
pred_server = dr.PredictionServer.list()[0]
deployment = dr.Deployment.create_from_custom_model_version(
    model_version.id,
    label=CUSTOM_MODEL_NAME,
    default_prediction_server_id=pred_server.id,
    max_wait=60 * 30,
)

print(f"Deployment ID: {deployment.id}")

import json

from datarobot_predict.deployment import predict_unstructured

# deployment = dr.Deployment.get("ADD_VALUE_HERE")
content, response_headers = predict_unstructured(
    deployment=deployment,
    data={
        "question": "How do I replace a custom model on an existing custom environment?",
        "filter": {"category": {"$eq": "index"}},
        "k": 1,
    },
)

json.loads(content)

use_case_id = "ADD_VALUE_HERE"
use_case = dr.UseCase.get(use_case_id)
# UNCOMMENT if you want to create a new Use Case
# use_case = dr.UseCase.create()

external_vdb_validation = CustomModelVectorDatabaseValidation.create(
    prompt_column_name="question", 
    target_column_name="relevant",
    deployment_id=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 Vector Database",
    use_case=use_case,
    validation_id=external_vdb_validation.id
)

assert vdb.execution_status == "COMPLETED"

print(f"Vector Database ID: {vdb.id}")