Transformers (ONNX) Embeddings

The TransformersEmbeddingModel is an EmbeddingModel implementation that locally computes sentence embeddings using a selected sentence transformer.spring-doc.cn

You can use any HuggingFace Embedding model.spring-doc.cn

It uses pre-trained transformer models, serialized into the Open Neural Network Exchange (ONNX) format.spring-doc.cn

The Deep Java Library and the Microsoft ONNX Java Runtime libraries are applied to run the ONNX models and compute the embeddings in Java.spring-doc.cn

Prerequisites

To run things in Java, we need to serialize the Tokenizer and the Transformer Model into ONNX format.spring-doc.cn

Serialize with optimum-cli - One, quick, way to achieve this, is to use the optimum-cli command line tool. The following snippet prepares a python virtual environment, installs the required packages and serializes (e.g. exports) the specified model using optimum-cli :spring-doc.cn

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

The snippet exports the sentence-transformers/all-MiniLM-L6-v2 transformer into the onnx-output-folder folder. The latter includes the tokenizer.json and model.onnx files used by the embedding model.spring-doc.cn

In place of the all-MiniLM-L6-v2 you can pick any huggingface transformer identifier or provide direct file path.spring-doc.cn

Auto-configuration

Spring AI provides Spring Boot auto-configuration for the ONNX Transformer Embedding Model. To enable it add the following dependency to your project’s Maven pom.xml file:spring-doc.cn

<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-transformers-spring-boot-starter</artifactId>
</dependency>

or to your Gradle build.gradle build file.spring-doc.cn

dependencies {
    implementation 'org.springframework.ai:spring-ai-transformers-spring-boot-starter'
}
Refer to the Dependency Management section to add the Spring AI BOM to your build file. Refer to the Repositories section to add these repositories to your build system.

To configure it, use the spring.ai.embedding.transformer.* properties.spring-doc.cn

For example, add this to your application.properties file to configure the client with the intfloat/e5-small-v2 text embedding model:spring-doc.cn

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

The complete list of supported properties are:spring-doc.cn

Embedding Properties

Property Description Default

spring.ai.embedding.transformer.enabledspring-doc.cn

Enable the Transformer Embedding model.spring-doc.cn

truespring-doc.cn

spring.ai.embedding.transformer.tokenizer.urispring-doc.cn

URI of a pre-trained HuggingFaceTokenizer created by the ONNX engine (e.g. tokenizer.json).spring-doc.cn

onnx/all-MiniLM-L6-v2/tokenizer.jsonspring-doc.cn

spring.ai.embedding.transformer.tokenizer.optionsspring-doc.cn

HuggingFaceTokenizer options such as ‘addSpecialTokens’, ‘modelMaxLength’, ‘truncation’, ‘padding’, ‘maxLength’, ‘stride’, ‘padToMultipleOf’. Leave empty to fallback to the defaults.spring-doc.cn

emptyspring-doc.cn

spring.ai.embedding.transformer.cache.enabledspring-doc.cn

Enable remote Resource caching.spring-doc.cn

truespring-doc.cn

spring.ai.embedding.transformer.cache.directoryspring-doc.cn

Directory path to cache remote resources, such as the ONNX modelsspring-doc.cn

${java.io.tmpdir}/spring-ai-onnx-modelspring-doc.cn

spring.ai.embedding.transformer.onnx.modelUrispring-doc.cn

Existing, pre-trained ONNX model.spring-doc.cn

onnx/all-MiniLM-L6-v2/model.onnxspring-doc.cn

spring.ai.embedding.transformer.onnx.modelOutputNamespring-doc.cn

The ONNX model’s output node name, which we’ll use for embedding calculation.spring-doc.cn

last_hidden_statespring-doc.cn

spring.ai.embedding.transformer.onnx.gpuDeviceIdspring-doc.cn

The GPU device ID to execute on. Only applicable if >= 0. Ignored otherwise.(Requires additional onnxruntime_gpu dependency)spring-doc.cn

-1spring-doc.cn

spring.ai.embedding.transformer.metadataModespring-doc.cn

Specifies what parts of the Documents content and metadata will be used for computing the embeddings.spring-doc.cn

NONEspring-doc.cn

Errors and special cases

If you see an error like Caused by: ai.onnxruntime.OrtException: Supplied array is ragged,.., you need to also enable the tokenizer padding in application.properties as follows:spring-doc.cn

spring.ai.embedding.transformer.tokenizer.options.padding=true

If you get an error like The generative output names don’t contain expected: last_hidden_state. Consider one of the available model outputs: token_embeddings, …​., you need to set the model output name to a correct value per your models. Consider the names listed in the error message. For example:spring-doc.cn

spring.ai.embedding.transformer.onnx.modelOutputName=token_embeddings

If you get an error like 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, that means that you model is larger than 2GB and is serialized in two files: model.onnx and model.onnx_data.spring-doc.cn

The model.onnx_data is called External Data and is expected to be under the same directory of the model.onnx.spring-doc.cn

Currently the only workaround is to copy the large model.onnx_data in the folder you run your Boot applicaiton.spring-doc.cn

If you get an error like ai.onnxruntime.OrtException: Error code - ORT_EP_FAIL - message: Failed to find CUDA shared provider, that means that you are using the GPU parameters spring.ai.embedding.transformer.onnx.gpuDeviceId , but the onnxruntime_gpu dependency are missing.spring-doc.cn

<dependency>
    <groupId>com.microsoft.onnxruntime</groupId>
    <artifactId>onnxruntime_gpu</artifactId>
</dependency>

Please select the appropriate onnxruntime_gpu version based on the CUDA version(ONNX Java Runtime).spring-doc.cn

Manual Configuration

If you are not using Spring Boot, you can manually configure the Onnx Transformers Embedding Model. For this add the spring-ai-transformers dependency to your project’s Maven pom.xml file:spring-doc.cn

<dependency>
  <groupId>org.springframework.ai</groupId>
  <artifactId>spring-ai-transformers</artifactId>
</dependency>
Refer to the Dependency Management section to add the Spring AI BOM to your build file.

then create a new TransformersEmbeddingModel instance and use the setTokenizerResource(tokenizerJsonUri) and setModelResource(modelOnnxUri) methods to set the URIs of the exported tokenizer.json and model.onnx files. (classpath:, file: or https: URI schemas are supported).spring-doc.cn

If the model is not explicitly set, TransformersEmbeddingModel defaults to sentence-transformers/all-MiniLM-L6-v2:spring-doc.cn

Dimensionsspring-doc.cn

384spring-doc.cn

Avg. performancespring-doc.cn

58.80spring-doc.cn

Speedspring-doc.cn

14200 sentences/secspring-doc.cn

Sizespring-doc.cn

80MBspring-doc.cn

The following snippet illustrates how to use the TransformersEmbeddingModel manually:spring-doc.cn

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"));
If you create an instance of TransformersEmbeddingModel manually, you must call the afterPropertiesSet() method after setting the properties and before using the client.

The first embed() call downloads the large ONNX model and caches it on the local file system. Therefore, the first call might take longer than usual. Use the #setResourceCacheDirectory(<path>) method to set the local folder where the ONNX models as stored. The default cache folder is ${java.io.tmpdir}/spring-ai-onnx-model.spring-doc.cn

It is more convenient (and preferred) to create the TransformersEmbeddingModel as a Bean. Then you don’t have to call the afterPropertiesSet() manually.spring-doc.cn

@Bean
public EmbeddingModel embeddingModel() {
   return new TransformersEmbeddingModel();
}