RustでVectorDBを触りたい？！
こんにちは、エンジニアの渡辺（@mochi_neko_7）です。
今回は Rust で VectorDB を触ってみる話をします。
VectorDB は機械学習の文脈で使われることが多いため Python のサポートをしているものが多い一方で、Rust で直接触れるものが本当に少ないです。
なぜわざわざ Rust で使用したいのかの背景もお話ししつつ、Rust で触れる VectorDB である Qdrant、 追加で必要になる Sentence Embeddings など実装する上で考慮すべきポイントもサンプルコードを交えながら紹介します。
かなりニッチなテーマになりますが LLM の文脈で Rust でサーバーを構築したい場合などにヒントになるかもしれません。
環境
本記事の検証で使用した環境は下記になります。
- Windows 11
- NVIDIA GeForce RTX 4090
- NVIDIA Driver 536.23
- CUDA 11.8.0
- cuDNN 8
- libtorch 2.0.1
- Docker Desktop 4.19.0
- Rust 1.71.1
- Qdrant 1.4.1
- rust-bert 0.21.0
背景
VectorDB（ベクトルデータベース）とは何かの説明は下記をご覧ください。
生成AI などの機械学習の文脈ではテキストや画像などのデータをベクトルに変換し、 ベクトル同士の類似度（距離や角度）を比較して検索することに特化したデータベースを利用します。
VectorDB を提供しているサービスは OSS も含めて様々ありますので、用途に合わせて選択することになります。
冒頭でも話したように VectorDB は機械学習の文脈で使われることが多いため、 ほとんどの VectorDB が Python で手軽に触れるよう整備されています。
例えば In-memory でサクッと使えるものだと ChromaDB などがあります。
ただし Python はプロトタイピングには困らない一方で実際のプロダクトとして複数人での保守運用も考えると、 Python ではなく Golang や Rust などの静的型付けのある言語で制約をつけながら開発をしたいという考え方もあります。
Python での使用例は既に数多くありますので、今回は後者、特に Rust を使用することを考えてみます。
ちなみに後で紹介する rust-bert という Python の Transformers のようなライブラリが出ていたり、 HuggingFace も Rust 向けの ML ライブラリ を開発していたりと Rust で ML の流れは少し垣間見えます。
VectorDB: Qdrant
2023/08/22 現在では Rust で触れる VectorDB は調べても下記くらいしか見つかりません。
Qdrant は本体も Rust 製ですね。
他の選択肢として、REST API や gRPC などの WebAPI が提供されている VectorDB なら クライアントを自分で実装すれば使用することは可能です。
今回は Qdrant を使用して検証してみます。
Sentence Embeddings: rust-bert
いざ Qdrant を使ってみようとドキュメントや API を見てみると、 データベースへのデータの登録はベクトルを生で受け取る API になっていることが分かります。
汎用的な VectorDB としては確かに十分ですが、 現実的にはテキストや画像などのデータをベクトル（Embeddings）に変換する必要があります。
他の VectorDB、例えば先に紹介した ChromaDB だとテキストの Embedding 処理が組み込みで提供されていますが、 どうやら Qdrant では自前で実装する必要があるようです。
公式のチュートリアルには Mightly というサービスを利用する例が紹介されています。
ただし Mightly は商用利用は有償のサービスのようなので注意が必要です。
OpenAI の Embeddings API を利用する手もありますが、こちらも従量課金制です。
他に選択肢がないか調べていたところ、次の記事で rust-bert という Rust で Transformer ベースの機械学習モデルが触れるライブラリがあることを知りました。
examples を覗いてみると Sentence Embeddings のサンプルがあるじゃないですか。
今回は rust-bert を使用してテキストをベクトルに変換する Sentence Embeddings をローカルで走らせてみましょう。
VectorDB 利用の流れ
今回の検証で行うことを整理します。
- Rust で VectorDB を利用する
- VectorDB は Qdrant を、クライアントは公式のものを利用する
- テキストデータを扱う
- Sentence Embeddings は rust-bert で HuggingFace のモデルを使用する
VectorDB の基本的な操作は次の２つです。
- データの登録（upsert）
- データの検索（search）
1. データの登録（upsert）
登録したいテキストデータを Sentence Embeddings モデルを利用してベクトルに変換し、 メタデータ（payload）と一緒にデータベースに登録（upsert）します。
2. データの検索（search）
検索したいテキスト（query）を Sentence Embeddings モデルを利用してベクトルに変換し、 そのベクトルと類似しているデータを検索（search）します。
結果は Vec<ScoredPoint> というスコア付きのデータで受け取ることができ、
メタデータ（payload）も含まれています。
メタデータは自由に定義することができるので、例えば元のテキストを残しておいたり、 日時などの付加情報を載せることもできます。
サンプルコード
概念の説明だけだと味気ないのでサンプルコードも見ていきましょう。
1. 前提
サンプルコードは下記の前提で動くことを想定しています。
- NVIDIA 製 GPU が利用できること
- NVIDIA Driver がインストールされていること（CUDA 11.8 以上）
- Docker が利用できること
これらのセットアップ手順は割愛します。
2. 環境構築
環境構築は Docker / Docker Compose で行っています。
rust-vector-db-demo/compose.yaml at main · mochi-neko/rust-vector-db-demo · GitHub
Rust の環境と別に、VectorDB の Qdrant のサービスを立てています。
Rust の環境は CUDA が使えるよう nvidia/cuda の 公式 Image をベースに、
Rust と libtorch を追加でインストールします。
rust-vector-db-demo/Docker/rust/Dockerfile at main · mochi-neko/rust-vector-db-demo · GitHub
CUDA、cuDNN、libtorch のバージョンは rust-bert が 利用している tch-rs の要件に合わせる必要がありますが、 今回は他で環境構築をしていたものを流用した関係で
- CUDA 11.8
- cuDNN 8
- libtorch 2.0.1
を使用してセットアップしています。
Docker Compose では GPU の利用のために deploy: のブロックの設定が必要になる点にご注意ください。
3. 利用する crate
Rust では下記の crate を使用しています。
[dependencies]
anyhow = "1.0.72"
chrono = "0.4.26"
qdrant-client = "1.4.0"
rust-bert = "0.21.0"
tokio = { version = "1.32.0", features = ["rt-multi-thread"] }
uuid = "1.4.1"
rust-vector-db-demo/Cargo.toml at main · mochi-neko/rust-vector-db-demo · GitHub
Rust に慣れていない方向けに用途だけ簡単に説明します。
- anyhow : エラーハンドリングライブラリ
- chrono : 日時のライブラリ
- qdrant-client : Qdrant の Rust クライアントライブラリ
- rust-bert : Transformer ベースの自然言語処理ライブラリ
- tokio : 非同期処理ライブラリ
- uuid : UUID ライブラリ（VectorDB の ID 発行のため）
プロダクションだと tracing などの Logging ライブラリも使用しますが、今回は省略しています。
以降のソースコードは main.rs にまとめて記述しています。
rust-vector-db-demo/src/main.rs at main · mochi-neko/rust-vector-db-demo · GitHub
4. Sentence Embeddings のモデルのセットアップ
Sentence Embeddings のモデルのセットアップは rust-bert のサンプルを参考に下記のコードで行っています。
// Setup sentence embeddings model // NOTE: Run blocking operation in task::spawn_blocking let model = task::spawn_blocking(move || { SentenceEmbeddingsBuilder::remote( SentenceEmbeddingsModelType::AllMiniLmL6V2, ) .create_model() }) .await??; // Get embedding dimension let dimension = model.get_embedding_dim()? as u64;
使用するモデルの種類は SentenceEmbeddingsModelType の enum から選択できます。
tokio を使って非同期 main 関数内でモデルのセットアップをする場合は、tokio::task::spawn_blocking() を使って
create_model() の内部のブロッキング操作を非同期のタスクとして分離しないと
実行時に panic になることに注意が必要です。
実際にはここでモデルのダウンロード処理が入るため、時間がかかります。
Qdrant の Collection 作成時にベクトルの次元数を指定する必要があるので、get_embedding_dim() で取得しておきます。
5. Qdrant のセットアップ
次に Qdrant のクライアントと Colletion をセットアップします。
// Setup Qdrant client let qdrant = QdrantClient::from_url("http://qdrant:6334").build()?; qdrant.health_check().await?; // Create collection let collection_name = "collection_name".to_string(); qdrant .delete_collection(collection_name.clone()) .await?; qdrant .create_collection(&CreateCollection { collection_name: collection_name.clone(), vectors_config: Some(VectorsConfig { config: Some(Config::Params(VectorParams { size: dimension, distance: Distance::Cosine.into(), ..Default::default() })), }), ..Default::default() }) .await?;
Qdrant の URL http://qdrant:6334 のホスト名は今回は Docker Compose のサービス名 qdrant で指定しています。
その後、Collection を作成しています。
Colletion はデータベースのテーブルのような概念です。
size: dimension, でベクトルの次元数を、
distance: Distance::Cosine.into(), で距離の判定ロジックをコサインに指定しています。
6. VectorDB の基本操作の実装
具体的な処理をする前に、VectorDB の基本操作の実装を確認します。
テキストをベクトル（Embeddings）に変換するには、SentenceEmbeddingsModel.encode() を使用します。
fn embed( model: &SentenceEmbeddingsModel, sentence: &String, ) -> Result<Vec<Vec<f32>>> { Ok(model.encode(&[sentence])?) }
データ登録の操作では、テキストをベクトルに変換してから PointStruct に整形してデータを登録します。
その際、Payload というデータにメタデータを載せることができるため、ここでは下記のデータを追加しています。
- 元のテキスト
- 作成した日時
async fn upsert( client: &QdrantClient, collection_name: &str, model: &SentenceEmbeddingsModel, text: &String, ) -> Result<()> { let embedding = embed(model, text)?; let mut points = Vec::new(); let mut payload: HashMap<String, Value> = HashMap::new(); payload.insert( "text".to_string(), Value::from(text.clone()), ); payload.insert( "datetime".to_string(), Value::from( chrono::Utc::now() .format("%Y-%m-%dT%H:%M:%S%.3f") .to_string(), ), ); for vector in embedding { let point = PointStruct::new( uuid::Uuid::new_v4().to_string(), vector, Payload::new_from_hashmap(payload.clone()), ); points.push(point); } client .upsert_points( collection_name.to_string(), points, None, ) .await?; Ok(()) }
データ検索の操作では、検索したいテキスト query をベクトルに変換してから search_points の API を叩いています。
オプションの count_limit: u64, で検索数の最大数を、filter: Option<Filter>, でフィルタを指定できます。
フィルタは例えば Payload に属性や Tag、日時などを仕込んで SQL の WHERE のように使用することもできます。
async fn search( client: &QdrantClient, collection_name: &str, model: &SentenceEmbeddingsModel, query: String, count_limit: u64, filter: Option<Filter>, ) -> Result<Vec<ScoredPoint>> { let embedding = embed(model, &query)?; let vector = embedding[0].clone(); let result = client .search_points(&SearchPoints { collection_name: collection_name.to_string(), vector, limit: count_limit, filter, with_payload: Some(true.into()), with_vectors: Some(true.into()), ..Default::default() }) .await?; Ok(result.result) }
7. サンプルの実行
Sentence-Transformers の Semantic Search のサンプルを借りて、実際に VectorDB を通して検索ができるか確認します。
まずはコーパスを VectorDB に登録します。
// Store corpus in Vector DB let corpus = vec![ "A man is eating food.", "A man is eating a piece of bread.", "The girl is carrying a baby.", "A man is riding a horse.", "A woman is playing violin.", "Two men pushed carts through the woods.", "A man is riding a white horse on an enclosed ground.", "A monkey is playing drums.", "A cheetah is running behind its prey.", ]; for sentence in corpus { upsert( &qdrant, &collection_name, &model, &sentence.to_string(), ) .await?; println!("Upserted: {}", sentence) }
近い意味のテキストを指定して検索をします。
// Search for similar sentences let queries = vec![ "A man is eating pasta.", "Someone in a gorilla costume is playing a set of drums.", "A cheetah chases prey on across a field.", ]; for query in queries { let result = search( &qdrant, &collection_name, &model, query.to_string(), 5, None, ) .await?; println!("Query: {}", query); for point in result { println!( "Score: {}, Text: {}", point.score, point .payload .get("text") .unwrap() ); } println!(); }
8. 実行結果
Docker の rust サービス内の bash で下記コマンドを叩いて main.rs を実行します。
cargo run
コマンドラインに表示される結果は下記になりました。
Upserted: A man is eating food. Upserted: A man is eating a piece of bread. Upserted: The girl is carrying a baby. Upserted: A man is riding a horse. Upserted: A woman is playing violin. Upserted: Two men pushed carts through the woods. Upserted: A man is riding a white horse on an enclosed ground. Upserted: A monkey is playing drums. Upserted: A cheetah is running behind its prey. Query: A man is eating pasta. Score: 0.70354867, Text: "A man is eating food." Score: 0.52719873, Text: "A man is eating a piece of bread." Score: 0.18889551, Text: "A man is riding a horse." Score: 0.10469921, Text: "A man is riding a white horse on an enclosed ground." Score: 0.09803035, Text: "A cheetah is running behind its prey." Query: Someone in a gorilla costume is playing a set of drums. Score: 0.64325327, Text: "A monkey is playing drums." Score: 0.25641555, Text: "A woman is playing violin." Score: 0.13887261, Text: "A man is riding a horse." Score: 0.11909151, Text: "A man is riding a white horse on an enclosed ground." Score: 0.10798677, Text: "A cheetah is running behind its prey." Query: A cheetah chases prey on across a field. Score: 0.8253214, Text: "A cheetah is running behind its prey." Score: 0.13989523, Text: "A man is eating food." Score: 0.12919357, Text: "A monkey is playing drums." Score: 0.10974166, Text: "A man is riding a white horse on an enclosed ground." Score: 0.06497804, Text: "A man is riding a horse."
近い意味のテキストがスコア順に並んでいることが分かり、動作確認ができました。
おわりに
以上のサンプルコードで Rust でも Qdrant （VectorDB）と rust-bert（Sentence Embeddings）を使うことで ローカルで VectorDB を触ることができました。
有料の他サービスを活用すればもっと手軽に実装することもできるかと思います。
VectorDB の具体的な用途はやはり LLM 関連が多いと思いますが、 AI Agent や AITuber のサーバーを Rust で実装することも要件次第では可能です。
WebAPI サーバーとして機能を実装するなら REST だと axum、 gRPC なら tonic などと組み合わせると良いです。
rust-bert の依存関係の問題でついでに tch-rs も動かせる環境がセットアップできているので、 例えば Stable Diffusion など他の機械学習関連のモデルも Rust で動かすことができたりします。
もちろん Python だともっと簡単に実装できるものですが、それ以外の言語の選択肢があるということも大切かもしれません。
Rust 人口がまだ少なく、かつ Rust で機械学習、LLM、VectorDB を触りたいという人はおそらく極端に少ないためかなりニッチな話になると思いますが、 個人的には Python より Rust の方が開発も保守運用もしやすいと思っているので要件の合う機会がありそうなら現場でも活用できないかなと思っています。
