集成文本搜索器

文本搜索允许在语料库中搜索语义上相似的文本。它通过将搜索查询嵌入到表示查询语义含义的高维向量中来工作，然后使用 ScaNN（可扩展最近邻）在预定义的自定义索引中进行相似性搜索。

与文本分类（例如 Bert 自然语言分类器）不同，扩展可以识别的项目数量不需要重新训练整个模型。新项目可以通过简单地重建索引来添加。这也使您可以处理更大的（100k+ 项目）语料库。

使用任务库 TextSearcher API 将您的自定义文本搜索器部署到您的移动应用程序中。

TextSearcher API 的主要功能

  • 将单个字符串作为输入，在索引中执行嵌入提取和最近邻搜索。

  • 输入文本处理，包括在图内或图外 Wordpiece 或 Sentencepiece 对输入文本进行分词。

先决条件

在使用 TextSearcher API 之前，需要根据要搜索的自定义文本语料库构建索引。这可以通过使用 Model Maker Searcher API 来实现，方法是遵循并调整 教程。

为此，您将需要

  • 一个 TFLite 文本嵌入器模型，例如通用句子编码器。例如，
    • 在 此 Colab 中重新训练的 模型，它针对设备上推理进行了优化。在 Pixel 6 上查询文本字符串只需 6 毫秒。
    • 在 此 Colab 中重新训练的 模型，它针对设备上推理进行了优化。在 Pixel 6 上查询文本字符串只需 6 毫秒。
  • 您的文本语料库。

完成此步骤后，您应该拥有一个独立的 TFLite 搜索器模型（例如 mobilenet_v3_searcher.tflite），它是在 TFLite 模型元数据 中附加了索引的原始文本嵌入器模型。

在 Java 中运行推理

步骤 1：导入 Gradle 依赖项和其他设置

将 .tflite 搜索器模型文件复制到将运行模型的 Android 模块的 assets 目录中。指定该文件不应压缩，并将 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 (NNAPI is included)
    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 Support Pypi 包。

您可以使用以下命令安装 TensorFlow Lite Support 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 演示工具。