Python 文本嵌入指南

借助 MediaPipe Text Embedder 任务,您可以创建文本数据的数字表示形式,以捕获其语义含义。以下说明介绍了如何在 Python 中使用文本嵌入器。

如需详细了解此任务的功能、模型和配置选项,请参阅概览

代码示例

Text Embedder 的示例代码在 Python 中提供了此任务的完整实现,以供您参考。此代码可帮助您测试此任务并开始构建自己的文本嵌入器。您可以通过 Google Colab 只使用网络浏览器查看、运行和修改文本嵌入器示例代码。您可以在 GitHub 上查看此示例的源代码。

初始设置

本部分介绍了专门针对使用文本嵌入器而设置开发环境和代码项目的关键步骤。如需了解如何为使用 MediaPipe 任务设置开发环境的一般信息(包括平台版本要求),请参阅 Python 设置指南

软件包

文本嵌入器使用 mediapipe pip 软件包。您可以安装以下依赖项:

$ python -m pip install mediapipe

导入

导入以下类来访问文本嵌入器任务函数:

import mediapipe as mp
from mediapipe.tasks import python
from mediapipe.tasks.python import text

模型

MediaPipe Text Embedder 任务需要一个与此任务兼容的经过训练的模型。如需详细了解文本嵌入器可用的经过训练的模型,请参阅任务概览“模型”部分

选择并下载模型,然后将其存储在本地目录中。您可以使用推荐的 UniversalSentenceEncoder 模型。

model_path = '/absolute/path/to/universal_sentence_encoder.tflite'

model_asset_path 参数中指定模型的路径,如下所示:

base_options = BaseOptions(model_asset_path=model_path)

创建任务

MediaPipe Text Embedder 任务使用 create_from_options 函数来设置该任务。create_from_options 函数接受配置选项的值,以设置嵌入器选项。您还可以使用 create_from_model_path 工厂函数初始化任务。create_from_model_path 函数接受经过训练的模型文件的相对或绝对路径。如需详细了解配置选项,请参阅配置选项

以下代码演示了如何构建和配置此任务。

import mediapipe as mp

BaseOptions = mp.tasks.BaseOptions
TextEmbedder = mp.tasks.text.TextEmbedder
TextEmbedderOptions = mp.tasks.text.TextEmbedderOptions

# For creating a text embedder instance:
options = TextEmbedderOptions(
    base_options=BaseOptions(model_asset_path=model_path),
    quantize=True)
text_embedder = TextEmbedder.create_from_options(options)

配置选项

此任务为 Python 应用提供以下配置选项:

选项名称 说明 值范围 默认值
l2_normalize 是否以 L2 范数归一化返回的特征向量。 仅当模型尚未包含原生 L2_NORMALIZATION TFLite 运算时,才可使用此选项。在大多数情况下,情况都是如此,L2 归一化是通过 TFLite 推断实现的,因此无需使用此选项。 Boolean False
quantize 是否应通过标量量化将返回的嵌入量化为字节。嵌入被隐式假定为单位范数,因此任何维度的值都必须在 [-1.0, 1.0] 之间。如果不属于这种情况,请使用 l2_normalize 选项。 Boolean False

准备数据

文本嵌入器可处理文本 (str) 数据。该任务会处理数据输入预处理,包括标记化和张量预处理。

所有预处理都在 embed 函数中处理。无需事先对输入文本进行额外的预处理。

input_text = "The input text to be embedded."

运行任务

Text Embedder 使用 embed 函数来触发推断。对于文本嵌入,这意味着返回输入文本的嵌入向量。

以下代码演示了如何使用任务模型执行处理。

# Perform text embedding on the provided input text.
embedding_result = text_embedder.embed(input_text)

处理和显示结果

文本嵌入器会输出一个 TextEmbedderResult,其中包含输入文本的嵌入列表(浮点数或标量量化)。

下面显示了此任务的输出数据示例:

TextEmbedderResult:
  Embedding #0 (sole embedding head):
    float_embedding: {0.2345f, 0.1234f, ..., 0.6789f}
    head_index: 0

您可以使用 TextEmbedder.cosine_similarity 函数比较两个嵌入的语义相似度。如需查看示例,请参阅以下代码。

# Compute cosine similarity.
similarity = TextEmbedder.cosine_similarity(
  embedding_result.embeddings[0],
  other_embedding_result.embeddings[0])