Transformers (ONNX) 嵌入
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
该代码片段将 sentence-transformers/all-MiniLM-L6-v2 转换器导出到 onnx-output-folder
文件夹中。后者包括嵌入模型使用的 tokenizer.json
和 model.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>
或者到你的 Gradle build.gradle
构建文件中。
dependencies {
implementation 'org.springframework.ai:spring-ai-transformers-spring-boot-starter'
}
要配置它,请使用 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.options | HuggingFaceTokenizer 的选项,如 ‘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.modelOutputName | ONNX 模型的输出节点名称,我们将使用它来计算嵌入。 | 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.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.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>
:::提示
请参考 依赖管理 部分,将 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();
}