跳到主要内容

Transformers (ONNX) 嵌入

DeepSeek V3 中英对照 (ONNX) Transformers Transformers (ONNX) Embeddings

TransformersEmbeddingModel 是一个 EmbeddingModel 实现,它使用选定的 sentence transformer 在本地计算 句子嵌入

你可以使用任何 HuggingFace Embedding 模型

它使用了预训练的 transformer 模型,并将其序列化为 Open Neural Network Exchange (ONNX) 格式。

Deep Java Library 和 Microsoft ONNX Java Runtime 库被应用于在 Java 中运行 ONNX 模型并计算嵌入。

先决条件

要在 Java 中运行程序,我们需要将 Tokenizer 和 Transformer 模型 序列化为 ONNX 格式。

使用 optimum-cli 进行序列化 - 一种快速实现这一目标的方法是使用 optimum-cli 命令行工具。以下代码片段准备了一个 Python 虚拟环境,安装了所需的包,并使用 optimum-cli 序列化(例如导出)指定的模型:

python3 -m venv venv
source ./venv/bin/activate
(venv) pip install --upgrade pip
(venv) pip install optimum onnx onnxruntime sentence-transformers
(venv) optimum-cli export onnx --model sentence-transformers/all-MiniLM-L6-v2 onnx-output-folder
bash

该代码片段将 sentence-transformers/all-MiniLM-L6-v2 转换器导出到 onnx-output-folder 文件夹中。后者包括嵌入模型使用的 tokenizer.jsonmodel.onnx 文件。

你可以选择任何 huggingface transformer 标识符或直接提供文件路径来替代 all-MiniLM-L6-v2

自动配置

Spring AI 为 ONNX Transformer Embedding Model 提供了 Spring Boot 自动配置。要启用它,请将以下依赖项添加到项目的 Maven pom.xml 文件中:

<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-transformers-spring-boot-starter</artifactId>
</dependency>
xml

或者到你的 Gradle build.gradle 构建文件中。

dependencies {
implementation 'org.springframework.ai:spring-ai-transformers-spring-boot-starter'
}
groovy
提示

请参考依赖管理部分,将 Spring AI BOM 添加到您的构建文件中。参考仓库部分,将这些仓库添加到您的构建系统中。

要配置它,请使用 spring.ai.embedding.transformer.* 属性。

例如,将此添加到您的 application.properties 文件中,以使用 intfloat/e5-small-v2 文本嵌入模型配置客户端:

spring.ai.embedding.transformer.onnx.modelUri=https://huggingface.co/intfloat/e5-small-v2/resolve/main/model.onnx
spring.ai.embedding.transformer.tokenizer.uri=https://huggingface.co/intfloat/e5-small-v2/raw/main/tokenizer.json

支持的完整属性列表如下:

嵌入属性

属性描述默认值
spring.ai.embedding.transformer.enabled启用 Transformer Embedding 模型。true
spring.ai.embedding.transformer.tokenizer.uri由 ONNX 引擎创建的预训练 HuggingFaceTokenizer 的 URI(例如 tokenizer.json)。onnx/all-MiniLM-L6-v2/tokenizer.json
spring.ai.embedding.transformer.tokenizer.optionsHuggingFaceTokenizer 的选项,如 ‘addSpecialTokens’, ‘modelMaxLength’, ‘truncation’, ‘padding’, ‘maxLength’, ‘stride’, ‘padToMultipleOf’。留空以回退到默认值。
spring.ai.embedding.transformer.cache.enabled启用远程资源缓存。true
spring.ai.embedding.transformer.cache.directory缓存远程资源的目录路径,例如 ONNX 模型。${java.io.tmpdir}/spring-ai-onnx-model
spring.ai.embedding.transformer.onnx.modelUri现有的预训练 ONNX 模型。onnx/all-MiniLM-L6-v2/model.onnx
spring.ai.embedding.transformer.onnx.modelOutputNameONNX 模型的输出节点名称,我们将使用它来计算嵌入。last_hidden_state
spring.ai.embedding.transformer.onnx.gpuDeviceId要执行的 GPU 设备 ID。仅当 >= 0 时适用。否则忽略。(需要额外的 onnxruntime_gpu 依赖)-1
spring.ai.embedding.transformer.metadataMode指定文档内容和元数据的哪些部分将用于计算嵌入。NONE

错误和特殊情况

备注

如果你看到类似 Caused by: ai.onnxruntime.OrtException: Supplied array is ragged,.. 的错误,你需要在 application.properties 中启用 tokenizer 的填充选项,如下所示:

spring.ai.embedding.transformer.tokenizer.options.padding=true
备注

如果你遇到类似 The generative output names don’t contain expected: last_hidden_state. Consider one of the available model outputs: token_embeddings, …​. 的错误,你需要根据你的模型将模型输出名称设置为正确的值。请参考错误消息中列出的名称。例如:

spring.ai.embedding.transformer.onnx.modelOutputName=token_embeddings
备注

如果你遇到类似 ai.onnxruntime.OrtException: Error code - ORT_FAIL - message: Deserialize tensor onnx::MatMul_10319 failed.GetFileLength for ./model.onnx_data failed:Invalid fd was supplied: -1 的错误,这意味着你的模型大于 2GB,并且被序列化为两个文件:model.onnxmodel.onnx_data

model.onnx_data 被称为 External Data,并且应该与 model.onnx 位于同一目录下。

目前唯一的解决方法是将较大的 model.onnx_data 文件复制到你运行 Boot 应用程序的文件夹中。

备注

如果你遇到类似 ai.onnxruntime.OrtException: Error code - ORT_EP_FAIL - message: Failed to find CUDA shared provider 的错误,这意味着你正在使用 GPU 参数 spring.ai.embedding.transformer.onnx.gpuDeviceId,但缺少 onnxruntime_gpu 依赖。

<dependency>
<groupId>com.microsoft.onnxruntime</groupId>
<artifactId>onnxruntime_gpu</artifactId>
</dependency>

请根据 CUDA 版本选择合适的 onnxruntime_gpu 版本(ONNX Java Runtime)。

手动配置

如果你没有使用 Spring Boot,你可以手动配置 Onnx Transformers 嵌入模型。为此,请将 spring-ai-transformers 依赖项添加到项目的 Maven pom.xml 文件中:

<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-transformers</artifactId>
</dependency>
xml

:::提示
请参考 依赖管理 部分,将 Spring AI BOM 添加到您的构建文件中。
:::

然后创建一个新的 TransformersEmbeddingModel 实例,并使用 setTokenizerResource(tokenizerJsonUri)setModelResource(modelOnnxUri) 方法来设置导出的 tokenizer.jsonmodel.onnx 文件的 URI。(支持 classpath:file:https: URI 模式)。

如果模型未明确设置,TransformersEmbeddingModel 默认使用 sentence-transformers/all-MiniLM-L6-v2

维度384
平均性能58.80
速度14200 句/秒
大小80MB

以下代码片段展示了如何手动使用 TransformersEmbeddingModel

TransformersEmbeddingModel embeddingModel = new TransformersEmbeddingModel();

// (optional) defaults to classpath:/onnx/all-MiniLM-L6-v2/tokenizer.json
embeddingModel.setTokenizerResource("classpath:/onnx/all-MiniLM-L6-v2/tokenizer.json");

// (optional) defaults to classpath:/onnx/all-MiniLM-L6-v2/model.onnx
embeddingModel.setModelResource("classpath:/onnx/all-MiniLM-L6-v2/model.onnx");

// (optional) defaults to ${java.io.tmpdir}/spring-ai-onnx-model
// Only the http/https resources are cached by default.
embeddingModel.setResourceCacheDirectory("/tmp/onnx-zoo");

// (optional) Set the tokenizer padding if you see an errors like:
// "ai.onnxruntime.OrtException: Supplied array is ragged, ..."
embeddingModel.setTokenizerOptions(Map.of("padding", "true"));

embeddingModel.afterPropertiesSet();

List<List<Double>> embeddings = this.embeddingModel.embed(List.of("Hello world", "World is big"));
java
备注

如果你手动创建 TransformersEmbeddingModel 的实例,必须在设置属性之后、使用客户端之前调用 afterPropertiesSet() 方法。

第一个 embed() 调用会下载大型的 ONNX 模型并将其缓存到本地文件系统中。因此,第一次调用可能会比平常花费更长的时间。你可以使用 #setResourceCacheDirectory(<path>) 方法来设置存储 ONNX 模型的本地文件夹。默认的缓存文件夹是 ${java.io.tmpdir}/spring-ai-onnx-model

TransformersEmbeddingModel 创建为一个 Bean 更为方便(且推荐)。这样你就不需要手动调用 afterPropertiesSet() 方法了。

@Bean
public EmbeddingModel embeddingModel() {
return new TransformersEmbeddingModel();
}
java