集成文本嵌入器。

文本嵌入器可将文本嵌入到表示其语义含义的高维特征向量中，然后将该向量与其他文本的特征向量进行比较，以评估它们的语义相似度。

与文本搜索不同，文本嵌入器允许即时计算文本之间的相似度，而不是搜索通过语料库构建的预定义索引。

使用 Task Library TextEmbedder API 将自定义文本嵌入器部署到移动应用中。

TextEmbedder API 的主要功能

支持的文本嵌入器模型

以下型号保证与 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 支持 Pypi 软件包。

您可以使用以下命令安装 TensorFlow Lite 支持 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 表示两个向量完全相同。

Cosine similarity: 0.954312

使用您自己的模型和测试数据试用简单的 TextEmbedder CLI 演示工具。

模型兼容性要求

TextEmbedder API 需要具有强制性 TFLite 模型元数据的 TFLite 模型。

支持三种主要类型的模型：

  • 基于 BERT 的模型（如需了解详情，请参阅源代码）：

    • 正好 3 个输入张量 (kTfLiteString)

    • ID 张量，元数据名称为“ids”，

    • 具有元数据名称“mask”的遮罩张量。

    • 细分受众群 ID 张量，元数据名称为“segment_ids”

    • 一个输出张量（kTfLiteUInt8/kTfLiteFloat32）

    • 具有 N 个组成部分，对应于此输出层的返回特征向量的 N 个维度。

    • 2 个或 4 个维度，即 [1 x N] 或 [1 x 1 x 1 x N]。

    • Wordpiece/Sentencepiece 词元化器的 input_process_units

  • 基于 Universal Sentence Encoder 的模型（如需了解详情，请参阅源代码）：

    • 正好 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)

    • 至少一个输出嵌入张量 (kTfLiteUInt8/kTfLiteFloat32)

    • 具有 N 个组成部分，对应于此输出层的返回特征向量的 N 个维度。

    • 2 个或 4 个维度，即 [1 x N] 或 [1 x 1 x 1 x N]。