Transformers (ONNX) 嵌入

TransformersEmbeddingModel 是一个 EmbeddingModel 实现，它使用选定的 句子变换器 在本地计算 句子嵌入。

你可以使用任何 HuggingFace 嵌入模型。

它使用 预训练 的变换器模型，序列化为 开放神经网络交换（ONNX） 格式。

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

前置条件

要在 Java 中运行，我们需要把 分词器（Tokenizer）和变换器模型序列化为 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

该片段把 sentence-transformers/all-MiniLM-L6-v2 变换器导出到 onnx-output-folder 文件夹。后者包含嵌入模型所使用的 tokenizer.json 和 model.onnx 文件。

你可以用任何 huggingface 变换器标识符或直接给出文件路径来替代 all-MiniLM-L6-v2。

自动配置

NOTE

Spring AI 的自动配置、starter 模块的制品名称发生了重大变更。 更多信息请参阅 升级说明。

Spring AI 为 ONNX 变换器嵌入模型提供 Spring Boot 自动配置。 要启用它，请将以下依赖添加到你的项目的 Maven pom.xml 文件：

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

或添加到你的 Gradle build.gradle 构建文件。

dependencies {
    implementation 'org.springframework.ai:spring-ai-starter-model-transformers'
}

TIP

参见 依赖管理 一节，把 Spring AI BOM 添加到你的构建文件。

参见 制品仓库 一节，把这些仓库添加到你的构建系统。

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

例如，把以下内容添加到你的 application.properties 文件，以用 intfloat/e5-small-v2 文本嵌入模型配置客户端：

spring.ai.embedding.transformer.onnx.model-uri=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

受支持属性的完整列表为：

嵌入属性

NOTE

嵌入自动配置的启用与禁用，现已通过以 spring.ai.model.embedding 为前缀的顶层属性来配置。

要启用，spring.ai.model.embedding=transformers（默认启用）

要禁用，spring.ai.model.embedding=none（或任何不匹配 transformers 的值）

做出此变更是为了支持配置多个模型。

属性	描述	默认值
spring.ai.embedding.transformer.enabled（已移除，不再生效）	启用变换器嵌入模型。	true
spring.ai.model.embedding	启用变换器嵌入模型。	transformers
spring.ai.embedding.transformer.tokenizer.uri	由 ONNX 引擎创建的预训练 HuggingFaceTokenizer 的 URI（例如 tokenizer.json）。	onnx/all-MiniLM-L6-v2/tokenizer.json
spring.ai.embedding.transformer.tokenizer.options	HuggingFaceTokenizer 选项，如 'addSpecialTokens'、'modelMaxLength'、'truncation'、'padding'、'maxLength'、'stride'、'padToMultipleOf'。留空则回退到默认值。	empty
spring.ai.embedding.transformer.cache.enabled	启用远程 Resource 缓存。	true
spring.ai.embedding.transformer.cache.directory	缓存诸如 ONNX 模型等远程资源的目录路径	${java.io.tmpdir}/spring-ai-onnx-model
spring.ai.embedding.transformer.onnx.model-uri	已有的、预训练的 ONNX 模型。	onnx/all-MiniLM-L6-v2/model.onnx
spring.ai.embedding.transformer.onnx.model-output-name	ONNX 模型的输出节点名称，用于嵌入计算。	lasthiddenstate
spring.ai.embedding.transformer.onnx.gpu-device-id	要执行所在的 GPU 设备 ID。仅在 >= 0 时适用，否则忽略。（需要额外的 onnxruntime_gpu 依赖）	-1
spring.ai.embedding.transformer.metadata-mode	指定文档的哪些内容和元数据部分将用于计算嵌入。	NONE

错误与特殊情况

NOTE

如果你看到类似 Caused by: ai.onnxruntime.OrtException: Supplied array is ragged,.. 的错误，你还需要在 application.properties 中启用分词器 padding，如下：

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

NOTE

如果你得到类似 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.model-output-name=token_embeddings

NOTE

如果你得到类似 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.onnx 和 model.onnx_data。

model.onnx_data 被称为 外部数据（External Data），预期与 model.onnx 位于同一目录下。

目前唯一的变通办法是把较大的 model.onnx_data 复制到你运行 Boot 应用的文件夹中。

NOTE

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

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

请根据 CUDA 版本选择适当的 onnxruntime_gpu 版本（ONNX Java Runtime）。

手动配置

如果你不使用 Spring Boot，可以手动配置 Onnx 变换器嵌入模型。 为此，将 spring-ai-transformers 依赖添加到你的项目的 Maven pom.xml 文件：

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

TIP

参见 依赖管理 一节，把 Spring AI BOM 添加到你的构建文件。

然后创建一个新的 TransformersEmbeddingModel 实例，并使用 setTokenizerResource(tokenizerJsonUri) 和 setModelResource(modelOnnxUri) 方法设置导出的 tokenizer.json 和 model.onnx 文件的 URI。（支持 classpath:、file: 或 https: URI 方案）。

若未显式设置模型，TransformersEmbeddingModel 默认使用 sentence-transformers/all-MiniLM-L6-v2：

维度	384
平均性能	58.80
速度	14200 sentences/sec
大小	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"));

NOTE

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

首次 embed() 调用会下载大型 ONNX 模型并缓存到本地文件系统。 因此，首次调用可能比平时耗时更长。 使用 #setResourceCacheDirectory(<path>) 方法设置存放 ONNX 模型的本地文件夹。 默认缓存文件夹为 ${java.io.tmpdir}/spring-ai-onnx-model。

把 TransformersEmbeddingModel 作为一个 Bean 来创建更便捷（也更推荐）。 那样你就不必手动调用 afterPropertiesSet()。

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