Python で TiDB + AI を使い始める

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

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

注記

TiDB Vector Search は現在ベータ版であり、 TiDB Cloudサーバーレスクラスターでのみ使用できます。

前提条件

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

• Python 3.8以上インストールされました。
• ギットインストールされました。
• TiDB Cloud Serverless クラスター。TiDB Cloud クラスターがない場合は、 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 Cloudの Vector Search と対話するための Python クライアント。
• sentence-transformers : テキストからベクトル埋め込みを生成するための事前トレーニング済みモデルを提供する Python ライブラリ。

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

1. クラスターページに移動し、ターゲット クラスターの名前をクリックして概要ページに移動します。

2. 右上隅の「接続」をクリックします。接続ダイアログが表示されます。

3. 接続ダイアログの構成が動作環境と一致していることを確認します。

   • 接続タイプはPublicに設定されています。

   • ブランチはmainに設定されています。

   • Connect With はSQLAlchemyに設定されています。

   • オペレーティング システムは環境に適合します。

   ヒント：

   プログラムが Windows Subsystem for Linux (WSL) で実行されている場合は、対応する Linux ディストリビューションに切り替えます。

4. PyMySQLタブをクリックし、接続文字列をコピーします。

   ヒント：

   まだパスワードを設定していない場合は、 「パスワードの生成」をクリックしてランダムなパスワードを生成します。

5. Python プロジェクトのルート ディレクトリに.envファイルを作成し、その中に接続文字列を貼り付けます。

   以下は macOS の例です。

   TIDB_DATABASE_URL="mysql+pymysql://<prefix>.root:<password>@gateway01.<region>.prod.aws.tidbcloud.com:4000/test?ssl_ca=/etc/ssl/cert.pem&ssl_verify_cert=true&ssl_verify_identity=true"
   

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

埋め込みモデルはデータをベクトル埋め込みに変換します。この例では、テキスト埋め込みに事前トレーニング済みのモデルmsmarco-MiniLM-L12-cos-v5を使用します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 次元のベクトルを生成します。

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 table which 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,
   # Determine whether to 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

出力から判断すると、泳いでいる動物は魚、または泳ぐ才能のある犬である可能性が高いです。

このデモでは、ベクトル検索が最も関連性の高いドキュメントを効率的に見つける方法を示します。検索結果はベクトルの近さによって整理されます。距離が小さいほど、ドキュメントの関連性が高くなります。

参照

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