Transformers (ONNX) 嵌入

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

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

应用 深度 Java 库 和 Microsoft ONNX Java 运行时 库来运行 ONNX 模型并在 Java 中计算嵌入。

序列化分词器和 Transformer 模型

为了在 Java 中运行，我们需要将分词器和 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
(venv) optimum-cli export onnx --generative sentence-transformers/all-MiniLM-L6-v2 onnx-output-folder

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

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

使用 ONNX 变换器模型

将 spring-ai-transformers 项目添加到您的 Maven 依赖项中。

<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 = 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();
}

Transformers Embedding Spring Boot 启动器

您可以使用以下 Spring Boot 启动器引导和自动装配 TransformersEmbeddingModel

<dependency>
   <groupId>org.springframework.ai</groupId>
   <artifactId>spring-ai-transformers-spring-boot-starter</artifactId>
</dependency>
请参阅 依赖项管理 部分，将 Spring AI BOM 添加到您的构建文件中。

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

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

spring.ai.embedding.transformer.onnx.modelUri=https://hugging-face.cn/intfloat/e5-small-v2/resolve/main/model.onnx
spring.ai.embedding.transformer.tokenizer.uri=https://hugging-face.cn/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”。留空以回退到默认值。

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.modelUri

现有的预训练 ONNX 模型。

onnx/all-MiniLM-L6-v2/model.onnx

spring.ai.embedding.transformer.onnx.gpuDeviceId

要执行的 GPU 设备 ID。仅当 >= 0 时适用。否则忽略。

-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