Transformers(ONNX)嵌入
TransformersEmbeddingModel 是 EmbeddingModel 的一种实现,它使用选定的 sentence transformer 在本地计算 sentence embeddings。
你可以使用任何 HuggingFace Embedding 模型。
它使用了预训练的 transformer 模型,并序列化为 Open Neural Network Exchange (ONNX) 格式。
Deep Java Library 和微软的 ONNX Java Runtime 库用于在 Java 中运行 ONNX 模型并计算嵌入向量。
前提条件
要在 Java 中运行,我们需要将 Tokenizer 和 Transformer Model 序列化为 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 文件。
除了 all-MiniLM-L6-v2,您可以选择任何 HuggingFace 转换器标识符,或提供直接的文件路径。
自动配置
Spring AI 的自动配置(auto-configuration)以及 starter 模块的构件名称(artifact names)发生了重大变更。详细信息请参阅 升级说明。
Spring AI 为 ONNX Transformer Embedding Model 提供了 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'
}
要对其进行配置,请使用 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
支持的属性完整列表如下:
嵌入属性
嵌入(embedding)自动配置项的启用与禁用现在通过前缀为 spring.ai.model.embedding 的顶层属性进行配置。
要启用,设置 spring.ai.model.embedding=transformers(默认已启用)。
要禁用,设置 spring.ai.model.embedding=none(或任何不等于 transformers 的值)。
此变更旨在支持多模型的配置。
| 属性 | 描述 | 默认值 |
|---|---|---|
| spring.ai.embedding.transformer.enabled (已移除,不再有效) | 启用 Transformer Embedding 模型。 | true |
| spring.ai.model.embedding | 启用 Transformer 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 | 启用远程资源缓存。 | 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 模型的输出节点名称。 | last_hidden_state |
| spring.ai.embedding.transformer.onnx.gpu-device-id | 要执行计算的 GPU 设备 ID。仅当值 >= 0 时有效,否则将被忽略。(需要额外的 onnxruntime_gpu 依赖项) | -1 |
| spring.ai.embedding.transformer.metadata-mode | 指定计算嵌入时将使用 Document 的内容和元数据的哪些部分。 | NONE |
错误与特殊情况
如果您看到类似 Caused by: ai.onnxruntime.OrtException: Supplied array is ragged,.. 的错误,您还需要在 application.properties 中启用分词器填充,如下所示:
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.model-output-name=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.onnx 和 model.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.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 Transformers Embedding Model。为此,请在项目的 Maven pom.xml 文件中添加 spring-ai-transformers 依赖:
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-transformers</artifactId>
</dependency>
参考 Dependency Management 章节,将 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 句/秒 |
| 大小 | 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"));
如果你手动创建 TransformersEmbeddingModel 的实例,在设置完属性后且在使用客户端之前,必须调用 afterPropertiesSet() 方法。
首次调用 embed() 会下载大型 ONNX 模型并将其缓存到本地文件系统。因此,第一次调用可能需要比平时更长的时间。可以使用 #setResourceCacheDirectory(<path>) 方法设置用于存储 ONNX 模型的本地文件夹。默认缓存文件夹为 ${java.io.tmpdir}/spring-ai-onnx-model。
创建 TransformersEmbeddingModel 作为 Bean 更为方便(且是首选方式)。这样,您就无需手动调用 afterPropertiesSet() 方法。
@Bean
public EmbeddingModel embeddingModel() {
return new TransformersEmbeddingModel();
}