Flutterでテキストをベクトル化する

テキストの意味検索や類似度計算をモバイルアプリで実現しようとしたとき、真っ先に思い浮かぶのはサーバーサイド API に投げる方法かと思います。しかし、ネットワーク遅延の排除・プライバシー保護・オフライン対応という観点では、オンデバイス推論は魅力的な選択肢です。

この記事では、BAAI/bge-small-en-v1.5 というテキスト埋め込みモデルを flutter_onnxruntime で動かし、ネットワーク接続なしでテキストをベクトルに変換するアプリを実装した内容を紹介します。

全体の構成

作ったものを大まかに整理するとこうなります。

Python 側でモデルを ONNX 形式に変換し、Flutter 側で純 Dart のトークナイザーと ONNX Runtime を使って推論するシンプルな構成です。

Python ツール：モデルを ONNX にエクスポートする

まず HuggingFace のモデルを ONNX 形式に変換する Python スクリプトを作ります。パッケージは uv で管理しています。

pyproject.toml

[project]
name = "vectorize-tools"
requires-python = ">=3.13"
dependencies = [
    "onnx>=1.20.1",
    "onnxruntime>=1.24.2",
    "onnxscript>=0.6.2",
    "torch>=2.10.0",
    "transformers>=5.2.0",
]

エクスポート自体は torch.onnx.export で行います。

def export_to_onnx(...) -> Path:
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModel.from_pretrained(model_name)
    model.eval()

    dummy_text = "This is a sample sentence."
    inputs = tokenizer(dummy_text, return_tensors="pt",
                       padding="max_length", max_length=128, truncation=True)

    with torch.no_grad():
        torch.onnx.export(
            model,
            args=(inputs["input_ids"], inputs["attention_mask"]),
            f=str(onnx_path),
            input_names=["input_ids", "attention_mask"],
            output_names=["last_hidden_state", "pooler_output"],
            dynamic_axes={
                "input_ids": {0: "batch_size", 1: "sequence_length"},
                "attention_mask": {0: "batch_size", 1: "sequence_length"},
                "last_hidden_state": {0: "batch_size", 1: "sequence_length"},
                "pooler_output": {0: "batch_size"},
            },
            opset_version=17,
        )

dynamic_axes を指定することで、バッチサイズやシーケンス長が推論時に可変になります。

IR バージョン問題

ここでハマりどころになるのが IR バージョンです。現在のコードベースでは Flutter 側に flutter_onnxruntime を使っており、Python 側も IR バージョン 10 を前提にモデルを再保存しています。エクスポートしたモデルの IR がそれより大きい場合は、アプリ側で読める値まで明示的に下げて保存し直します。

def save_single_file_model(model: onnx.ModelProto, path: Path) -> None:
    onnx.save_model(
        model,
        str(path),
        save_as_external_data=False,  # テンソルを .onnx に埋め込む
        all_tensors_to_one_file=True,
    )
    # 古い .onnx.data ファイルが残っていれば削除
    external_data_path = path.with_suffix(path.suffix + ".data")
    if external_data_path.exists():
        external_data_path.unlink()


# エクスポート後の処理
onnx_model = onnx.load(str(onnx_path))
if onnx_model.ir_version > target_ir_version:
    onnx_model.ir_version = target_ir_version
save_single_file_model(onnx_model, onnx_path)
onnx.checker.check_model(onnx_model)

torch.onnx.export が出力するファイルは .onnx と .onnx.data に分かれた外部データ形式になる場合があります。save_as_external_data=False で保存し直すことで単一の .onnx ファイルに統合でき、Flutter 側でのアセット管理がシンプルになります。

現在のスクリプトでは CLI から --target-ir-version を指定でき、既定値は 10 です。

uv run main.py --model BAAI/bge-small-en-v1.5 --output ../assets/models --target-ir-version 10

Flutter 側で Unsupported model IR version が出た場合は、まず生成済みモデルの IR version を確認し、tools/main.py と同じ 10 をターゲットにして再生成するのが前提になります。

最後に ORT セッションで推論テストをして、埋め込み次元と L2 ノルムを確認します。

session = ort.InferenceSession(str(onnx_path))
outputs = session.run(None, {
    "input_ids": inputs["input_ids"].numpy(),
    "attention_mask": inputs["attention_mask"].numpy(),
})
embedding = outputs[0][0].mean(axis=0)
print(f"  埋め込みベクトル次元: {embedding.shape[0]}")
print(f"  L2ノルム: {np.linalg.norm(embedding):.4f}")

これでモデルの準備は完了です。uv run main.py で実行すると assets/models/BAAI_bge-small-en-v1.5.onnx が生成されます。

Flutter：ONNX Runtime でモデルを動かす

パッケージの導入

pubspec.yaml に flutter_onnxruntime と path_provider を追加し、モデルと vocab をアセットとして登録します。

pubspec.yaml

dependencies:
  flutter_onnxruntime: ^1.6.3
  path_provider: ^2.1.5

flutter:
  assets:
    - assets/models/BAAI_bge-small-en-v1.5.onnx
    - assets/models/vocab.txt

モデルのロード

今回の実装では、モデル asset を一時ディレクトリへ毎回書き出してから createSession に渡しています。createSessionFromAsset も使えますが、内部キャッシュの都合で asset 更新が反映されにくいので、アプリ側で明示的に上書きしています。

Future<void> initialize() async {
  await _closeSession();

  final data = await rootBundle.load(modelAssetPath);
  final directory = await getTemporaryDirectory();
  final fileName = modelAssetPath.split('/').last;
  final modelFile = File('${directory.path}${Platform.pathSeparator}$fileName');
  await modelFile.writeAsBytes(
    data.buffer.asUint8List(data.offsetInBytes, data.lengthInBytes),
    flush: true,
  );
  final vocabText = await rootBundle.loadString(vocabAssetPath);

  _tokenizer = BertWordPieceTokenizer.fromVocabText(vocabText);
  _session = await OnnxRuntime().createSession(modelFile.path);
}

推論の実行

ONNX Runtime への入力は OrtValue として渡します。今回のモデルは input_ids / attention_mask ともに tensor(int64) を期待するので、Int64List を明示的に渡すのがポイントです。

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,
});

try {
  final outputValue = outputs['last_hidden_state']!;
  final tensorData = await outputValue.asList();
  final embedding = _meanPoolEmbedding(tensorData, encoded.attentionMask);
} finally {
  for (final output in outputs.values) {
    await output.dispose();
  }
  await inputTensor.dispose();
  await attentionTensor.dispose();
}

推論結果の last_hidden_state は [1, sequence_length, hidden_size] の形状で返ってきます。これを mean pooling してテキスト全体の埋め込みベクトルを得ます。

純 Dart で WordPiece Tokenizer を実装する

ONNX モデルに渡す input_ids を生成するには、テキストをトークンに分割する必要があります。Python なら transformers の AutoTokenizer をそのまま使えますが、Flutter でこれを呼び出す手段はありません。

そこで vocab.txt を読み込んで BERT の WordPiece アルゴリズムを純 Dart で実装しています。

vocab.txt の構造

[PAD]       ← 0 行目 = id:0
[unused1]
...
[UNK]       ← id:100
[CLS]       ← id:101
[SEP]       ← id:102
...
hello       ← 普通の単語
##ing       ← サブワード（先頭以外のピース）

行番号がそのままトークン ID になります。この構造を利用して Map<String, int> と List<String> の2つのテーブルを作ります。

factory BertWordPieceTokenizer.fromVocabText(String vocabText) {
  final lines = vocabText
      .split('\n')
      .map((line) => line.trim())
      .where((line) => line.isNotEmpty)
      .toList(growable: false);

  final tokenToId = <String, int>{};
  for (var i = 0; i < lines.length; i++) {
    tokenToId[lines[i]] = i;
  }
  // [CLS], [SEP], [UNK], [PAD] の id を取得して保持
  ...
}

WordPiece アルゴリズム

WordPiece は「最長一致で単語をサブワードに分割する」アルゴリズムです。たとえば vectorize という単語が vocab になくても、vector + ##ize のように分割できます。

List<int> _wordPieceTokenizeToIds(String token) {
  // 単語全体が vocab にあればそのまま使う
  final direct = _tokenToId[token];
  if (direct != null) return [direct];

  // 最長一致で左から順にサブワードを決定する
  final pieces = <int>[];
  var start = 0;
  while (start < token.length) {
    int? currentId;
    var end = token.length;
    while (start < end) {
      final sub = token.substring(start, end);
      final candidate = start == 0 ? sub : '##$sub';  // 先頭以外は ## を付ける
      final id = _tokenToId[candidate];
      if (id != null) {
        currentId = id;
        break;
      }
      end--;
    }
    if (currentId == null) return [unkId];  // 分割不能なら [UNK]
    pieces.add(currentId);
    start = end;
  }
  return pieces;
}

BERT 入力形式への変換

BERT 系モデルの入力は [CLS] トークン列 [SEP] という形式です。また、モデルに渡す配列長は固定（今回は 128）なので、足りない部分を [PAD] で埋めます。

({List<int> inputIds, List<int> attentionMask}) encode(String text, {int maxLength = 128}) {
  final normalized = text.toLowerCase();  // BGE は uncased モデル
  final basicTokens = _basicTokenize(normalized);

  final wordPieceIds = <int>[];
  for (final token in basicTokens) {
    wordPieceIds.addAll(_wordPieceTokenizeToIds(token));
  }

  // [CLS] + tokens + [SEP]
  final inputIds = <int>[clsId, ...wordPieceIds, sepId];

  // maxLength を超えた場合は末尾を切り詰め、最後を [SEP] に保つ
  if (inputIds.length > maxLength) {
    inputIds
      ..removeRange(maxLength - 1, inputIds.length)
      ..[maxLength - 1] = sepId;
  }

  // 実トークン:1、PAD:0 の attention mask
  final attentionMask = List<int>.filled(inputIds.length, 1, growable: true);
  while (inputIds.length < maxLength) {
    inputIds.add(padId);
    attentionMask.add(0);
  }

  return (inputIds: inputIds, attentionMask: attentionMask);
}

Mean Pooling：attention mask を使って正しく集約する

last_hidden_state はシーケンス中の全トークン分の出力を持っています。このうち [PAD] トークンの出力は意味を持たないため、attention_mask を使って除外した上で平均を取ることが重要です。

単純に全トークンを平均してしまうと、テキストが短いほど大量の PAD が混入し、埋め込みの品質が下がります。

static List<double> _meanPoolEmbedding(dynamic tensorValue, List<int> attentionMask) {
  final tokenVectors = _extractTokenVectors(tensorValue);
  final count = tokenVectors.length < attentionMask.length
      ? tokenVectors.length
      : attentionMask.length;
  final hiddenSize = tokenVectors.first.length;
  final sums = List<double>.filled(hiddenSize, 0);
  var validTokenCount = 0;

  for (var tokenIndex = 0; tokenIndex < count; tokenIndex++) {
    if (attentionMask[tokenIndex] == 0) continue;  // PAD はスキップ
    validTokenCount++;
    final tokenVector = tokenVectors[tokenIndex];
    for (var i = 0; i < hiddenSize; i++) {
      sums[i] += tokenVector[i];
    }
  }

  if (validTokenCount == 0) {
    throw StateError('No valid tokens for pooling.');
  }

  return sums.map((value) => value / validTokenCount).toList();
}

これを呼び出す側では、トークナイズの結果得た attentionMask をそのまま渡します。

final embedding = _meanPoolEmbedding(
  (firstTensor as OrtValueTensor).value,
  encoded.attentionMask,  // PAD 位置が 0 になっているマスク
);

実行画面

テキストを入力して「Vectorize」を押すと、トークン列と埋め込みベクトルの先頭16次元が表示されます。

まとめ

実装してみてわかったポイントをまとめておきます。

ポイント	内容
外部データ形式	エクスポート時に .onnx + .onnx.data に分かれることがある。save_as_external_data=False で単一ファイルに統合する
モデルロード	fromBuffer を使えばアセットからファイルコピーなしで直接ロードできる
Mean Pooling	attention_mask で PAD トークンを除外しないと短いテキストの精度が落ちる
モデルサイズ	ONNX モデルは数十MB〜100MB超になることがあり、アプリサイズ増加・配布コスト増・初回ロード時間の増加につながる

完全なソースコードは GitHub リポジトリを参照してください。

この先にできること

オンデバイスで埋め込みベクトルが得られると、さまざまな応用に発展できます。

意味検索
あらかじめ各ドキュメントを数値ベクトル（埋め込み）に変換してローカルに保存しておき、クエリのベクトルとのコサイン類似度で検索する意味検索を実現できます。キーワード一致ではなく、「意味が近い」文書を見つけられるのが強みです。

テキスト間の類似度計算
2つのテキストをそれぞれベクトル化し、コサイン類似度を計算するだけで、重複検出・推薦・クラスタリングなどに応用できます。

多言語モデルへの切り替え
BAAI/bge-m3 のような多言語対応モデルに差し替えると、日本語を含む 100 言語以上のテキストをベクトル化できます。Python ツールの --model オプションを変更するだけで試せます。

INT8 量子化によるモデルの軽量化
Python ツールには --quantize オプションがあります。INT8 量子化を適用するとモデルサイズを大幅に削減でき、モバイル端末での推論速度向上も期待できます。

最後に

株式会社ボトルキューブではFlutterを使ったお仕事を募集中です。
お問い合わせフォームからご連絡ください。

また、一緒に働く仲間も募集しています。
詳細は採用情報ページをご覧ください。