转换器 (ONNX) 嵌入
这TransformersEmbeddingModel
是一个EmbeddingModel
使用选定的 sentence transformer 在本地计算句子嵌入的实现。
您可以使用任何 HuggingFace Embedding 模型。
它使用预先训练的 transformer 模型,序列化为 Open Neural Network Exchange (ONNX) 格式。
Deep Java 库和 Microsoft ONNX Java 运行时库用于运行 ONNX 模型并计算 Java 中的嵌入。
先决条件
要在 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
嵌入模型使用的文件。
您可以选择任何 huggingface 转换器标识符或提供直接文件路径,而不是全 MiniLM-L6-v2。
自动配置
Spring AI 为 ONNX Transformer 嵌入模型提供了 Spring Boot 自动配置。
要启用它,请将以下依赖项添加到项目的 Maven 中pom.xml
文件:
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-transformers-spring-boot-starter</artifactId>
</dependency>
或发送到您的 Gradlebuild.gradle
build 文件。
dependencies {
implementation 'org.springframework.ai:spring-ai-transformers-spring-boot-starter'
}
请参阅 Dependency Management 部分,将 Spring AI BOM 添加到您的构建文件中。 请参阅 Repositories 部分,将这些存储库添加到您的构建系统中。 |
要配置它,请使用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 模型。 |
真 |
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 |
启用远程资源缓存。 |
真 |
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.model输出名称 |
ONNX 模型的输出节点名称,我们将用于嵌入计算。 |
last_hidden_state |
spring.ai.embedding.transformer.onnx.gpuDeviceId |
要执行的 GPU 设备 ID。仅当 >= 0 时适用。否则为 Ignored。(需要额外的 onnxruntime_gpu 依赖项) |
-1 |
spring.ai.embedding.transformer.metadataMode |
指定 Documents 内容和元数据的哪些部分将用于计算嵌入。 |
没有 |
错误和特殊情况
如果您看到类似 spring.ai.embedding.transformer.tokenizer.options.padding=true |
如果您收到类似 spring.ai.embedding.transformer.onnx.modelOutputName=token_embeddings |
如果您收到类似 这 目前,唯一的解决方法是将 |
如果您收到类似 <dependency> <groupId>com.microsoft.onnxruntime</groupId> <artifactId>onnxruntime_gpu</artifactId> </dependency> 请根据 CUDA 版本(ONNX Java Runtime)选择合适的 onnxruntime_gpu 版本。 |
手动配置
如果您不使用 Spring Boot,则可以手动配置 Onnx Transformers Embedding Model。
为此,请添加spring-ai-transformers
依赖项添加到项目的 Mavenpom.xml
文件:
<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
文件。(classpath:
,file:
或https:
支持 URI 架构)。
如果未显式设置模型,则TransformersEmbeddingModel
默认为 sentence-transformers/all-MiniLM-L6-v2:
尺寸 |
384 |
平均性能 |
58.80 |
速度 |
14200 句/秒 |
大小 |
80兆字节 |
以下代码段说明了如何使用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()
call 下载大型 ONNX 模型并将其缓存在本地文件系统上。
因此,第一次调用可能需要比平时更长的时间。
使用#setResourceCacheDirectory(<path>)
方法将 ONNX 模型的存储位置设置为本地文件夹。
默认缓存文件夹为${java.io.tmpdir}/spring-ai-onnx-model
.
将 TransformersEmbeddingModel 创建为Bean
.
然后,您不必调用afterPropertiesSet()
手动地。
@Bean
public EmbeddingModel embeddingModel() {
return new TransformersEmbeddingModel();
}