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

您可以使用任何HuggingFace Embedding模型

它使用预训练的转换器模型,序列化为开放神经网络交换 (ONNX) 格式。

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

先决条件

要在 Java 中运行,我们需要将分词器和转换器模型序列化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.jsonmodel.onnx 文件。

您可以使用任何 huggingface 转换器标识符或提供直接文件路径来代替 all-MiniLM-L6-v2。

自动配置

Spring AI 自动配置和启动器模块的工件名称已发生重大更改。 有关更多信息,请参阅升级说明

Spring AI 为 ONNX Transformer Embedding 模型提供 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 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

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

Embedding 属性

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.uriONNX 引擎创建的预训练 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 中启用分词器填充,如下所示:

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 称为外部数据,并且应位于 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 Embedding 模型。 为此,请将 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.jsonmodel.onnx 文件的 URI。(支持 classpath:file:https: URI 方案)。

如果未显式设置模型,TransformersEmbeddingModel 默认为 sentence-transformers/all-MiniLM-L6-v2

维度平均性能速度大小
38458.8014200 个句子/秒80MB

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

TransformersEmbeddingModel embeddingModel = new TransformersEmbeddingModel();

// (可选) 默认为 classpath:/onnx/all-MiniLM-L6-v2/tokenizer.json
embeddingModel.setTokenizerResource("classpath:/onnx/all-MiniLM-L6-v2/tokenizer.json");

// (可选) 默认为 classpath:/onnx/all-MiniLM-L6-v2/model.onnx
embeddingModel.setModelResource("classpath:/onnx/all-MiniLM-L6-v2/model.onnx");

// (可选) 默认为 ${java.io.tmpdir}/spring-ai-onnx-model
// 默认情况下仅缓存 http/https 资源。
embeddingModel.setResourceCacheDirectory("/tmp/onnx-zoo");

// (可选) 如果看到类似以下错误,请设置分词器填充:
// "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();
}

文档有误?请协助编辑

发现文档问题?点击此处直接在 GitHub 上编辑并提交 PR,帮助我们改进文档!