Transformers (ONNX) Embeddings

`TransformersEmbeddingClient`是 `EmbeddingClient`实现,它使用所选 sentence transformer在本地计算 sentence embeddings。 它使用 pre-trained 转换器模型,序列化成 Open Neural Network Exchange (ONNX) 格式。 Deep Java Library 和 Microsoft ONNX Java Runtime 库用于运行 ONNX 模型和计算 Java 中的嵌入。

Serialize the Tokenizer and the Transformer Model

要在 Java 中运行,我们需要将分词器和转换器模型序列化为 ONNX 格式。

Serialize with 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.jsonmodel.onnx

你可以选择任何 huggingface 变换器标识符来代替 all-MiniLM-L6-v2 或提供直接的文件路径。

Using the ONNX Transformers models

spring-ai-transformers 项目添加到你的 maven 依赖项:

<dependency>
  <groupId>org.springframework.ai</groupId>
  <artifactId>spring-ai-transformers</artifactId>
</dependency>
  1. 参见 Dependency Management 部分,将 Spring AI BOM 添加到你的构建文件中。

然后创建一个新的 TransformersEmbeddingClient 实例,并使用 setTokenizerResource(tokenizerJsonUri)setModelResource(modelOnnxUri) 方法设置导出 tokenizer.jsonmodel.onnx 文件的 URI。(支持`classpath:`, file:https: URI 架构)。

如果没有明确设置模型,则 TransformersEmbeddingClient 默认值为 ` sentence-transformers/all-MiniLM-L6-v2`:

Dimensions

384

Avg. performance

58.80

Speed

14200 sentences/sec

Size

80MB

以下代码段说明了如何手动使用 TransformersEmbeddingClient

TransformersEmbeddingClient embeddingClient = new TransformersEmbeddingClient();

// (optional) defaults to classpath:/onnx/all-MiniLM-L6-v2/tokenizer.json
embeddingClient.setTokenizerResource("classpath:/onnx/all-MiniLM-L6-v2/tokenizer.json");

// (optional) defaults to classpath:/onnx/all-MiniLM-L6-v2/model.onnx
embeddingClient.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.
embeddingClient.setResourceCacheDirectory("/tmp/onnx-zoo");

// (optional) Set the tokenizer padding if you see an errors like:
// "ai.onnxruntime.OrtException: Supplied array is ragged, ..."
embeddingClient.setTokenizerOptions(Map.of("padding", "true"));

embeddingClient.afterPropertiesSet();

List<List<Double>> embeddings = embeddingClient.embed(List.of("Hello world", "World is big"));

手动创建时,必须在设置属性并使用客户端之前调用 afterPropertiesSet()

第一个 embed() 调用下载大型 ONNX 模型并将其缓存在本地文件系统上。因此,第一次调用可能比平时花费的时间长。使用 #setResourceCacheDirectory(<path>) 方法设置存储 ONNX 模型的本地文件夹。默认缓存文件夹是 ${java.io.tmpdir}/spring-ai-onnx-model

更为方便(也是首选)的是将 TransformersEmbeddingClient 创建为一个 Bean。然后你就不必手动调用 afterPropertiesSet()

@Bean
public EmbeddingClient embeddingClient() {
   return new TransformersEmbeddingClient();
}

Transformers Embedding Spring Boot Starter

你可以使用以下 Spring Boot starter 启动并自动装配 TransformersEmbeddingClient

<dependency>
   <groupId>org.springframework.ai</groupId>
   <artifactId>spring-ai-transformers-spring-boot-starter</artifactId>
</dependency>
  1. 参见 Dependency Management 部分,将 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

受支持的属性的完整列表:

Property Description Default

spring.ai.embedding.transformer.enabled

启用变压器嵌入客户端。

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

Enable remote Resource caching.

true

spring.ai.embedding.transformer.cache.directory

要缓存诸如 ONNX 模型之类的远程资源的目录路径

${java.io.tmpdir}/spring-ai-onnx-model

spring.ai.embedding.transformer.onnx.modelUri

Existing, pre-trained ONNX model.

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