前回ONNX Runtimeを使ってテキストをベクトル化してみたので、今回はその続きとして、実際の検索機能を実装してみます。
実行画面
先にどのような画面になるか出しておきます。
ソースコードはGitHubで公開しています。
実装方針
検索することを目的として、できるだけ簡素に実装します。
ベクトル化したカタログデータ(検索対象)を用意し、Dart コードでコサイン類似度を全件走査します。
- カタログ(book / album 混在)を Python で事前ベクトル化
- L2 正規化済みインデックスを
assetsに同梱 - Flutter 起動時にインデックスを一括ロード
- クエリを ONNX で埋め込み化して top-10 類似検索
全体構成
1. カタログデータの用意
今回は検索対象のカタログを JSON で用意して、最低限のデータを対象にします。
{
"items": [
{
"id": "book-001",
"type": "book",
"title": "...",
"summary": "...",
"genre": "...",
"author": "..."
},
{
"id": "album-002",
"type": "album",
"title": "...",
"summary": "...",
"genre": "...",
"artist": "..."
}
]
}
ベクトル化するテキストは title だけでなく summary や genre を連結し、キーワード一致では拾えない文脈(雰囲気・テーマ)で検索できるようにします。
2. Python で検索インデックスを事前計算する
今回追加した tools/build_index.py は、カタログを読み込んで次の処理を行います。
- 各アイテムの埋め込み対象テキストを組み立てる
- ONNX Runtime で
last_hidden_stateを取得 - attention mask 付き mean pooling
- L2 正規化
-
search_index.jsonに保存
実行コマンド:
cd tools
uv run build_index.py \
--catalog ../assets/data/catalog.dummy.500.json \
--onnx-model ../assets/models/BAAI_bge-small-en-v1.5.onnx \
--tokenizer-model BAAI/bge-small-en-v1.5 \
--output ../assets/data/search_index.json
出力は以下の形式です。
{
"model": "BAAI_bge-small-en-v1.5.onnx",
"dimension": 384,
"items": [
{
"id": "book-001",
"title": "...",
"summary": "...",
"genre": "...",
"embedding": [0.012345, -0.004321, "..."]
}
]
}
embedding を事前に正規化しておくことで、Flutter 側は「クエリを正規化して内積を取るだけ」でコサイン類似度と同等の結果を得られます。
3. Flutter 側の読み込みと検索
3.1 インデックスのロード
SearchIndex.loadFromAsset(...) で search_index.json を読み込み、アプリ起動時にメモリ展開します。
final index = await SearchIndex.loadFromAsset('assets/data/search_index.json');
dimension と各 embedding.length を突き合わせることで、壊れたデータの混入を最小限で防げます。
3.2 検索文字列のベクトル化
検索文字列は VectorizerService の embed(String) でベクトル化します。現在の実装では flutter_onnxruntime を使って ONNX モデルを読み込んでいます。
final embedding = await _vectorizerService.embed(query);
VectorizerService 側では次の点を押さえています。
- ONNX モデル asset を一時ディレクトリへ書き出して
OnnxRuntime().createSession(...)に渡す -
input_ids/attention_maskをInt64Listで作り、モデルが要求するtensor(int64)に合わせる - 推論後に
OrtValueを毎回dispose()してネイティブメモリを解放する
イメージとしては以下のような処理です。
final inputTensor = await OrtValue.fromList(
Int64List.fromList(encoded.inputIds),
[1, maxLength],
);
final attentionTensor = await OrtValue.fromList(
Int64List.fromList(encoded.attentionMask),
[1, maxLength],
);
final outputs = await session.run({
'input_ids': inputTensor,
'attention_mask': attentionTensor,
});
final outputValue = outputs['last_hidden_state']!;
final tensorData = await outputValue.asList();
final embedding = _meanPoolEmbedding(tensorData, encoded.attentionMask);
3.3 検索処理
コサイン類似度を使って検索します。ベクトル化したテキストとカタログデータの埋め込みベクトルとの内積を取ります。
cosine_search.dart 側では、次の流れで top-k を返します。
- クエリベクトルを L2 正規化
- 全件に対して内積を計算
- スコア降順ソート
take(10)
final matches = topKCosineNormalized(
query: embedding,
index: index,
k: 10,
);
件数によっては Isolate を分ける、ANN ライブラリを使う、sqlite-vss のような別実装へ寄せるなど、最適化の余地があります。
4. ONNX モデル運用上の注意
現在の Python エクスポートスクリプト tools/main.py は、Flutter 側との互換性を保つために ONNX モデルの IR version を 10 に揃えています。
uv run main.py --model BAAI/bge-small-en-v1.5 --output ../assets/models --target-ir-version 10
この調整を入れているのは、エクスポート環境のライブラリ更新によって生成される IR version が変わることがあるためです。アプリ側で Unsupported model IR version が出た場合は、まずこの値を確認して再生成するのが前提になります。
5. 現時点の制約
今回のモデルは BAAI/bge-small-en-v1.5(英語中心)です。今回は検索すること自体を要件にしているので、日本語対応は後回しにしています。
多言語化する場合は、次の 2 点をセットで見直すのが実務上安全です。
- 多言語埋め込みモデルへ差し替え
- Flutter 側トークナイザー実装の互換性確認
モデルだけ変えてトークナイザーを据え置くと、意図どおりの入力にならないケースがあります。
この先にできること
フィールド重み付け検索
title を強め(重視する)、summary を弱める(軽視する)など、フィールドごとに重みを変えて、検索結果の傾向を調整する。この辺りは要件次第です。
再ランキング
埋め込み検索で top-50 を取ってから、軽量なルールベースや別モデルで再ランキングするなどの、検索アルゴリズムの調整。フィールドの重み付けと同様に要件次第です。
差分更新インデックス
カタログデータをネットワークからダウンロードするなどで、アプリサイズを減らしたり、カタログ更新を容易にする。
多言語移行
モデルとトークナイザーの組み合わせを再検証して、日本語を含む実データで評価指標(nDCG / Recall@k など)を回すと移行判断がしやすくなります。
高速化
今回は Dart で全件走査を実装しましたが、件数が増えてきたら Isolate を分けてUIのフリーズを防いだり、ANN ライブラリ(例: hnswlib の Dart バインディング)を導入するのも選択肢になります。
最後に
株式会社ボトルキューブではFlutterを使ったお仕事を募集中です。
お問い合わせフォームからご連絡ください。
また、一緒に働く仲間も募集しています。
詳細は採用情報ページをご覧ください。
