Python で TiDB + AI を使い始める

このチュートリアルでは、セマンティック検索機能を提供するシンプルな AI アプリケーションの開発方法を説明します。従来のキーワード検索とは異なり、セマンティック検索はクエリの背後にある意味をインテリジェントに理解し、最も関連性の高い結果を返します。たとえば、「犬」、「魚」、「木」というタイトルのドキュメントがあり、「泳ぐ動物」を検索すると、アプリケーションは「魚」を最も関連性の高い結果として識別します。

このチュートリアルでは、 TiDB ベクトル検索 、Python、 Python 用 TiDB ベクター SDK 、AI モデルを使用してこの AI アプリケーションを開発します。

注記：

ベクトル検索機能は、TiDB セルフマネージド クラスターとTiDB Cloudサーバーレスクラスターでのみ使用できます。

前提条件

このチュートリアルを完了するには、次のものが必要です。

• Python 3.8以上インストールされました。
• ギットインストールされました。
• TiDB クラスター。

TiDB クラスターがない場合は、次のように作成できます。

• ローカルテストTiDBクラスタをデプロイまたは本番のTiDBクラスタをデプロイに従ってローカル クラスターを作成します。
• TiDB Cloud Serverless クラスターの作成に従って、独自のTiDB Cloudクラスターを作成します。

始める

次の手順では、アプリケーションをゼロから開発する方法を示します。デモを直接実行するには、 pingcap/tidb-vector-pythonリポジトリのサンプル コードをチェックアウトできます。

ステップ1. 新しいPythonプロジェクトを作成する

任意のディレクトリに、新しい Python プロジェクトとexample.pyという名前のファイルを作成します。

mkdir python-client-quickstart
cd python-client-quickstart
touch example.py

ステップ2. 必要な依存関係をインストールする

プロジェクト ディレクトリで、次のコマンドを実行して必要なパッケージをインストールします。

pip install sqlalchemy pymysql sentence-transformers tidb-vector python-dotenv

• tidb-vector : TiDB ベクトル検索と対話するための Python クライアント。
• sentence-transformers : テキストからベクトル埋め込み生成するための事前トレーニング済みモデルを提供する Python ライブラリ。

ステップ3. TiDBクラスターへの接続文字列を構成する

選択した TiDB デプロイメント オプションに応じて、クラスター接続文字列を構成します。

• TiDB Cloud Serverless
• TiDB Self-Managed

TIDB_DATABASE_URL="mysql+pymysql://<USER>:<PASSWORD>@<HOST>:<PORT>/<DATABASE>"
# For example: TIDB_DATABASE_URL="mysql+pymysql://root@127.0.0.1:4000/test"

ステップ4. 埋め込みモデルを初期化する

埋め込みモデルデータをベクトル埋め込みに変換します。この例では、テキスト埋め込みに事前トレーニング済みのモデルmsmarco-MiniLM-L12-cos-v5を使用します。7 sentence-transformersによって提供されるこの軽量モデルは、テキスト データを 384 次元のベクトル埋め込みに変換します。

モデルを設定するには、次のコードをexample.pyファイルにコピーします。このコードはSentenceTransformerインスタンスを初期化し、後で使用するためにtext_to_embedding()関数を定義します。

from sentence_transformers import SentenceTransformer

print("Downloading and loading the embedding model...")
embed_model = SentenceTransformer("sentence-transformers/msmarco-MiniLM-L12-cos-v5", trust_remote_code=True)
embed_model_dims = embed_model.get_sentence_embedding_dimension()

def text_to_embedding(text):
    """Generates vector embeddings for the given text."""
    embedding = embed_model.encode(text)
    return embedding.tolist()

ステップ5. TiDBクラスターに接続する

TiDBVectorClientクラスを使用して TiDB クラスターに接続し、ベクター列を持つテーブルembedded_documentsを作成します。

注記

テーブル内のベクトル列の次元が、埋め込みモデルによって生成されたベクトルの次元と一致していることを確認します。たとえば、 msmarco-MiniLM-L12-cos-v5モデルは 384 次元のベクトルを生成するため、 embedded_documentsのベクトル列の次元も 384 である必要があります。

import os
from tidb_vector.integrations import TiDBVectorClient
from dotenv import load_dotenv

# Load the connection string from the .env file
load_dotenv()

vector_store = TiDBVectorClient(
   # The 'embedded_documents' table will store the vector data.
   table_name='embedded_documents',
   # The connection string to the TiDB cluster.
   connection_string=os.environ.get('TIDB_DATABASE_URL'),
   # The dimension of the vector generated by the embedding model.
   vector_dimension=embed_model_dims,
   # Recreate the table if it already exists.
   drop_existing_table=True,
)

ステップ6. テキストデータを埋め込み、ベクトルを保存する

このステップでは、「dog」、「fish」、「tree」などの単語を 1 つ含むサンプル ドキュメントを準備します。次のコードは、 text_to_embedding()関数を使用してこれらのテキスト ドキュメントをベクトル埋め込みに変換し、ベクトル ストアに挿入します。

documents = [
    {
        "id": "f8e7dee2-63b6-42f1-8b60-2d46710c1971",
        "text": "dog",
        "embedding": text_to_embedding("dog"),
        "metadata": {"category": "animal"},
    },
    {
        "id": "8dde1fbc-2522-4ca2-aedf-5dcb2966d1c6",
        "text": "fish",
        "embedding": text_to_embedding("fish"),
        "metadata": {"category": "animal"},
    },
    {
        "id": "e4991349-d00b-485c-a481-f61695f2b5ae",
        "text": "tree",
        "embedding": text_to_embedding("tree"),
        "metadata": {"category": "plant"},
    },
]

vector_store.insert(
    ids=[doc["id"] for doc in documents],
    texts=[doc["text"] for doc in documents],
    embeddings=[doc["embedding"] for doc in documents],
    metadatas=[doc["metadata"] for doc in documents],
)

ステップ7. セマンティック検索を実行する

このステップでは、既存のドキュメント内のどの単語とも直接一致しない「泳ぐ動物」を検索します。

次のコードでは、 text_to_embedding()関数を再度使用してクエリ テキストをベクトル埋め込みに変換し、その埋め込みを使用してクエリを実行して、最も近い 3 つの一致を検索します。

def print_result(query, result):
   print(f"Search result (\"{query}\"):")
   for r in result:
      print(f"- text: \"{r.document}\", distance: {r.distance}")

query = "a swimming animal"
query_embedding = text_to_embedding(query)
search_result = vector_store.query(query_embedding, k=3)
print_result(query, search_result)

example.pyファイルを実行すると、出力は次のようになります。

Search result ("a swimming animal"):
- text: "fish", distance: 0.4562914811223072
- text: "dog", distance: 0.6469335836410557
- text: "tree", distance: 0.798545178640937

検索結果の 3 つの用語は、クエリされたベクトルからのそれぞれの距離によって並べ替えられます。距離が小さいほど、対応するdocumentの関連性が高くなります。

したがって、出力によると、泳いでいる動物は魚、または泳ぐ才能のある犬である可能性が最も高いです。

参照

• ベクトルデータ型
• ベクター検索インデックス