テキスト埋め込みツールを統合する。

テキスト エンベッダーを使用すると、テキストをそのセマンティックな意味を表す高次元の特徴ベクトルに埋め込むことができます。このベクトルは、他のテキストの特徴ベクトルと比較して、セマンティックな類似性を評価できます。

テキスト検索とは異なり、テキスト エンベッダーを使用すると、コーパスから構築された事前定義済みのインデックスを検索するのではなく、テキスト間の類似性をその場で計算できます。

Task Library TextEmbedder API を使用して、カスタム テキスト エンベッダーをモバイルアプリにデプロイします。

TextEmbedder API の主な機能

  • 入力テキストの処理。入力テキストに対するグラフ内またはグラフ外の Wordpiece または Sentencepiece トークン化など。

  • 特徴ベクトル間のコサイン類似度を計算する組み込みのユーティリティ関数。

サポートされているテキスト エンベッダー モデル

次のモデルは、TextEmbedder API との互換性が保証されています。

C++ で推論を実行する

// Initialization.
TextEmbedderOptions options:
options.mutable_base_options()->mutable_model_file()->set_file_name(model_path);
std::unique_ptr<TextEmbedder> text_embedder = TextEmbedder::CreateFromOptions(options).value();

// Run inference with your two inputs, `input_text1` and `input_text2`.
const EmbeddingResult result_1 = text_embedder->Embed(input_text1);
const EmbeddingResult result_2 = text_embedder->Embed(input_text2);

// Compute cosine similarity.
double similarity = TextEmbedder::CosineSimilarity(
    result_1.embeddings[0].feature_vector()
    result_2.embeddings[0].feature_vector());

TextEmbedder を構成するその他のオプションについては、ソースコードをご覧ください。

Python で推論を実行する

ステップ 1: TensorFlow Lite Support Pypi パッケージをインストールします。

TensorFlow Lite Support Pypi パッケージは、次のコマンドを使用してインストールできます。

pip install tflite-support

ステップ 2: モデルを使用する

from tflite_support.task import text

# Initialization.
text_embedder = text.TextEmbedder.create_from_file(model_path)

# Run inference on two texts.
result_1 = text_embedder.embed(text_1)
result_2 = text_embedder.embed(text_2)

# Compute cosine similarity.
feature_vector_1 = result_1.embeddings[0].feature_vector
feature_vector_2 = result_2.embeddings[0].feature_vector
similarity = text_embedder.cosine_similarity(
    result_1.embeddings[0].feature_vector, result_2.embeddings[0].feature_vector)

TextEmbedder を構成するその他のオプションについては、ソースコードをご覧ください。

検索結果の例

正規化された特徴ベクトル間のコサイン類似度は、-1 から 1 までのスコアを返します。値が高いほど優れています。つまり、コサイン類似度が 1 の場合、2 つのベクトルは同一です。

Cosine similarity: 0.954312

独自のモデルとテストデータを使用して、シンプルな TextEmbedder 用 CLI デモツールを試す。

モデルの互換性要件

TextEmbedder API は、必須の TFLite モデル メタデータを含む TFLite モデルを想定しています。

主に次の 3 種類のモデルがサポートされています。

  • BERT ベースのモデル(詳しくは、ソースコードをご覧ください):

    • 3 つの入力テンソル(kTfLiteString)

    • メタデータ名「ids」の ID テンソル。

    • メタデータ名「mask」のマスク テンソル。

    • セグメント ID テンソル(メタデータ名「segment_ids」)

    • 出力テンソル(kTfLiteUInt8/kTfLiteFloat32)が 1 つだけ

    • この出力レイヤの返された特徴ベクトルの N ディメンションに対応する N コンポーネント。

    • 2 つまたは 4 つのディメンション([1 x N] または [1 x 1 x 1 x N])。

    • Wordpiece/Sentencepiece Tokenizer の input_process_units

  • ユニバーサル センテンス エンコーダ ベースのモデル(詳細については、ソースコードをご覧ください)。

    • 3 つの入力テンソル(kTfLiteString)

    • メタデータ名「inp_text」を含むクエリ テキスト テンソル。

    • 応答コンテキスト テンソル。メタデータ名は「res_context」です。

    • メタデータ名「res_text」を含むレスポンス テキスト テンソル。

    • 出力テンソルが 2 つ(kTfLiteUInt8/kTfLiteFloat32)

    • メタデータ名「query_encoding」を持つクエリ エンコード テンソル。

    • レスポンス エンコード テンソル。メタデータ名は「response_encoding」。

    • どちらも、この出力レイヤの返された特徴ベクトルの N ディメンションに対応する N コンポーネントを使用します。

    • どちらも 2 次元または 4 次元([1 x N] または [1 x 1 x 1 x N])です。

  • 次の条件を満たすテキスト エンベッダー モデル:

    • 入力テキスト テンソル(kTfLiteString)
    • 少なくとも 1 つの出力エンベディング テンソル(kTfLiteUInt8/kTfLiteFloat32)

    • この出力レイヤの返された特徴ベクトルの N ディメンションに対応する N コンポーネント。

    • 2 つまたは 4 つのディメンション([1 x N] または [1 x 1 x 1 x N])。