整合文字搜尋工具

文字搜尋功能可搜尋語料庫中語意相似的文字。這項功能會將搜尋查詢嵌入高維度向量,代表查詢的語意,然後使用 ScaNN (可擴充的最鄰近項目) 在預先定義的自訂索引中執行相似度搜尋。

與文字分類 (例如 Bert 自然語言分類器) 不同,擴充可辨識的項目數量不需要重新訓練整個模型。只要重新建立索引,即可新增項目。這也支援處理較大的語料庫 (超過 10 萬個項目)。

使用 Task Library TextSearcher API,將自訂文字搜尋器部署到行動應用程式中。

TextSearcher API 的主要功能

  • 以單一字串做為輸入內容,並在索引中執行嵌入擷取和最鄰近搜尋。

  • 處理輸入文字,包括在圖表內或圖表外對輸入文字進行 WordpieceSentencepiece 權杖化。

必要條件

使用 TextSearcher API 前,必須先根據要搜尋的自訂文字語料庫建立索引。您可以按照教學課程的說明,使用 Model Maker Searcher API 達成這個目標。

如要執行這項操作,您需要:

  • TFLite 文字嵌入模型,例如 Universal Sentence Encoder。例如:
    • 這個Colab 中重新訓練的模型,已針對裝置端推論進行最佳化。在 Pixel 6 上查詢文字字串只需 6 毫秒。
    • 量化模型,這個模型比上述模型小,但每個嵌入內容需要 38 毫秒。
  • 文字語料庫。

完成這個步驟後,您應該會取得獨立的 TFLite 搜尋器模型 (例如 mobilenet_v3_searcher.tflite),這是原始的文字嵌入模型,其中附加了索引到 TFLite 模型中繼資料

在 Java 中執行推論作業

步驟 1:匯入 Gradle 依附元件和其他設定

.tflite 搜尋器模型檔案複製到 Android 模組的資產目錄,模型將在該處執行。指定檔案不應壓縮,並將 TensorFlow Lite 程式庫新增至模組的 build.gradle 檔案:

android {
    // Other settings

    // Specify tflite index file should not be compressed for the app apk
    aaptOptions {
        noCompress "tflite"
    }

}

dependencies {
    // Other dependencies

    // Import the Task Vision Library dependency
    implementation 'org.tensorflow:tensorflow-lite-task-vision:0.4.4'
    // Import the GPU delegate plugin Library for GPU inference
    implementation 'org.tensorflow:tensorflow-lite-gpu-delegate-plugin:0.4.4'
}

步驟 2:使用模型

// Initialization
TextSearcherOptions options =
    TextSearcherOptions.builder()
        .setBaseOptions(BaseOptions.builder().useGpu().build())
        .setSearcherOptions(
            SearcherOptions.builder().setL2Normalize(true).build())
        .build();
TextSearcher textSearcher =
    textSearcher.createFromFileAndOptions(context, modelFile, options);

// Run inference
List<NearestNeighbor> results = textSearcher.search(text);

如要進一步瞭解如何設定 TextSearcher,請參閱原始碼和 Javadoc

以 C++ 執行推論

// Initialization
TextSearcherOptions options;
options.mutable_base_options()->mutable_model_file()->set_file_name(model_path);
options.mutable_embedding_options()->set_l2_normalize(true);
std::unique_ptr<TextSearcher> text_searcher = TextSearcher::CreateFromOptions(options).value();

// Run inference with your input, `input_text`.
const SearchResult result = text_searcher->Search(input_text).value();

如要進一步瞭解如何設定 TextSearcher,請參閱原始碼

在 Python 中執行推論

步驟 1:安裝 TensorFlow Lite 支援 Pypi 套件。

您可以使用下列指令安裝 TensorFlow Lite 支援 Pypi 套件:

pip install tflite-support

步驟 2:使用模型

from tflite_support.task import text

# Initialization
text_searcher = text.TextSearcher.create_from_file(model_path)

# Run inference
result = text_searcher.search(text)

如要設定 TextSearcher 的其他選項,請參閱原始碼

搜尋結果範例

Results:
 Rank#0:
  metadata: The sun was shining on that day.
  distance: 0.04618
 Rank#1:
  metadata: It was a sunny day.
  distance: 0.10856
 Rank#2:
  metadata: The weather was excellent.
  distance: 0.15223
 Rank#3:
  metadata: The cat is chasing after the mouse.
  distance: 0.34271
 Rank#4:
  metadata: He was very happy with his newly bought car.
  distance: 0.37703

使用自己的模型和測試資料,試用簡單的 TextSearcher CLI 示範工具