Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Improve ONNX Tansformer Embedding Model

Browse Source 
- Fix model output name handling.
 - Improve documentation.
This commit is contained in:
Christian Tzolov
2024-07-09 11:21:46 +02:00
parent f61ccdd09d
commit dd22342c0c
4 changed files with 102 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
models/spring-ai-transformers/src/main/java/org/springframework/ai/transformers/TransformersEmbeddingModel.java
View File

@@ -200,7 +200,9 @@ public class TransformersEmbeddingModel extends AbstractEmbeddingModel implement
logger.info("Model output names: " + onnxModelOutputs.stream().collect(Collectors.joining(", ")));
Assert.isTrue(onnxModelOutputs.contains(this.modelOutputName),
"The generative output names doesn't contain expected: " + this.modelOutputName);
"The generative output names doesn't contain expected: " + this.modelOutputName
+ ". Consider one of the available model outputs: "
+ onnxModelOutputs.stream().collect(Collectors.joining(", ")));
}
private Resource getCachedResource(Resource resource) {

2
pom.xml
View File

@@ -155,7 +155,7 @@
<jackson.version>2.16.1</jackson.version>
<djl.version>0.28.0</djl.version>
<onnxruntime.version>1.17.0</onnxruntime.version>
<onnxruntime.version>1.18.0</onnxruntime.version>
<com.google.cloud.version>26.41.0</com.google.cloud.version>
<qdrant.version>1.9.1</qdrant.version>
<spring-retry.version>2.0.5</spring-retry.version>

157
spring-ai-docs/src/main/antora/modules/ROOT/pages/api/embeddings/onnx.adoc
View File

@@ -2,18 +2,17 @@
The `TransformersEmbeddingModel` is an `EmbeddingModel` implementation that locally computes https://www.sbert.net/examples/applications/computing-embeddings/README.html#sentence-embeddings-with-transformers[sentence embeddings] using a selected https://www.sbert.net/[sentence transformer].
You can use any link:https://huggingface.co/spaces/mteb/leaderboard[HuggingFace Embedding model].
It uses https://www.sbert.net/docs/pretrained_models.html[pre-trained] transformer models, serialized into the https://onnx.ai/[Open Neural Network Exchange (ONNX)] format.
The https://djl.ai/[Deep Java Library] and the Microsoft https://onnxruntime.ai/docs/get-started/with-java.html[ONNX Java Runtime] libraries are applied to run the ONNX models and compute the embeddings in Java.
== Serialize the Tokenizer and the Transformer Model
== Prerequisites
To run things in Java, we need to serialize the Tokenizer and the Transformer Model into ONNX format.
=== Serialize with optimum-cli
One, quick, way to achieve this, is to use the https://huggingface.co/docs/optimum/exporters/onnx/usage_guides/export_a_model#exporting-a-model-to-onnx-using-the-cli[optimum-cli] command line tool.
To run things in Java, we need to *serialize the Tokenizer and the Transformer Model* into `ONNX` format.
Serialize with optimum-cli - One, quick, way to achieve this, is to use the https://huggingface.co/docs/optimum/exporters/onnx/usage_guides/export_a_model#exporting-a-model-to-onnx-using-the-cli[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` :
[source,bash]
@@ -21,17 +20,104 @@ The following snippet prepares a python virtual environment, installs the requir
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
(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 https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2[sentence-transformers/all-MiniLM-L6-v2] transformer into the `onnx-output-folder` folder. Later includes the `tokenizer.json` and `model.onnx` files used by the embedding model.
In place of the all-MiniLM-L6-v2 you can pick any huggingface transformer identifier or provide direct file path.
== Using the ONNX Transformers models
== Auto-configuration
Add the `spring-ai-transformers` project to your maven dependencies:
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:
[source, xml]
----
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-transformers-spring-boot-starter</artifactId>
</dependency>
----
or to your Gradle `build.gradle` build file.
[source,groovy]
----
dependencies {
implementation 'org.springframework.ai:spring-ai-transformers-spring-boot-starter'
}
----
TIP: Refer to the xref:getting-started.adoc#dependency-management[Dependency Management] section to add the Spring AI BOM to your build file.
Refer to the xref:getting-started.adoc#repositories[Repositories] section to add these repositories to your build system.
To configure it, use the `spring.ai.embedding.transformer.*` properties.
For example, add this to your _application.properties_ file to configure the client with the https://huggingface.co/intfloat/e5-small-v2[intfloat/e5-small-v2] text embedding model:
----
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:
=== Embedding Properties
[cols="3*"]
|===
| Property | Description | Default
| spring.ai.embedding.transformer.enabled | Enable the Transformer Embedding model. | true
| spring.ai.embedding.transformer.tokenizer.uri | URI of a pre-trained HuggingFaceTokenizer created by the ONNX engine (e.g. tokenizer.json). | onnx/all-MiniLM-L6-v2/tokenizer.json
| spring.ai.embedding.transformer.tokenizer.options | HuggingFaceTokenizer options such as '`addSpecialTokens`', '`modelMaxLength`', '`truncation`', '`padding`', '`maxLength`', '`stride`', '`padToMultipleOf`'. Leave empty to fallback to the defaults. | empty
| spring.ai.embedding.transformer.cache.enabled | Enable remote Resource caching. | true
| spring.ai.embedding.transformer.cache.directory | Directory path to cache remote resources, such as the ONNX models | ${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.modelOutputName | The ONNX model's output node name, which we'll use for embedding calculation. | last_hidden_state
| spring.ai.embedding.transformer.onnx.gpuDeviceId | The GPU device ID to execute on. Only applicable if >= 0. Ignored otherwise. | -1
| spring.ai.embedding.transformer.metadataMode | Specifies what parts of the Documents content and metadata will be used for computing the embeddings. | NONE
|===
=== Errors and special cases
[NOTE]
====
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.ai.embedding.transformer.tokenizer.options.padding=true
----
====
[NOTE]
====
If you get an error like `The generative output names doesn'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.
Cosider the names listed in the error message.
For example:
----
spring.ai.embedding.transformer.onnx.modelOutputName=token_embeddings
----
====
[NOTE]
====
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`.
The `model.onnx_data` is called link:https://onnx.ai/onnx/repo-docs/ExternalData.html#external-data[External Data] and is expected to be under the same directory of the `model.onnx`.
Currently the only workaround is to copy the large `model.onnx_data` in the folder you run your Boot applicaiton.
====
== 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:
[source,xml]
----
@@ -99,54 +185,3 @@ public EmbeddingModel embeddingModel() {
}
----
== Transformers Embedding Spring Boot Starter
You can bootstrap and autowire the `TransformersEmbeddingModel` with the following Spring Boot starter:
[source,xml]
----
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-transformers-spring-boot-starter</artifactId>
</dependency>
----
TIP: Refer to the xref:getting-started.adoc#dependency-management[Dependency Management] section to add the Spring AI BOM to your build file.
To configure it, use the `spring.ai.embedding.transformer.*` properties.
For example, add this to your _application.properties_ file to configure the client with the https://huggingface.co/intfloat/e5-small-v2[intfloat/e5-small-v2] text embedding model:
----
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:
[cols="3*"]
|===
| Property | Description | Default
| spring.ai.embedding.transformer.enabled | Enable the Transformer Embedding model. | true
| spring.ai.embedding.transformer.tokenizer.uri | URI of a pre-trained HuggingFaceTokenizer created by the ONNX engine (e.g. tokenizer.json). | onnx/all-MiniLM-L6-v2/tokenizer.json
| spring.ai.embedding.transformer.tokenizer.options | HuggingFaceTokenizer options such as '`addSpecialTokens`', '`modelMaxLength`', '`truncation`', '`padding`', '`maxLength`', '`stride`', '`padToMultipleOf`'. Leave empty to fallback to the defaults. | empty
| spring.ai.embedding.transformer.cache.enabled | Enable remote Resource caching. | true
| spring.ai.embedding.transformer.cache.directory | Directory path to cache remote resources, such as the ONNX models | ${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.modelOutputName | The ONNX model's output node name, which we'll use for embedding calculation. | last_hidden_state
| spring.ai.embedding.transformer.onnx.gpuDeviceId | The GPU device ID to execute on. Only applicable if >= 0. Ignored otherwise. | -1
| spring.ai.embedding.transformer.metadataMode | Specifies what parts of the Documents content and metadata will be used for computing the embeddings. | NONE
|===
NOTE: 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.ai.embedding.transformer.tokenizer.options.padding=true
----
NOTE: If you get an error like `Exception in thread "main" java.lang.IllegalArgumentException: The generative output names doesn't contain expected: last_hidden_state`, you need to set the model output name to a correct value per your models. For example:
----
spring.ai.embedding.transformer.onnx.modelOutputName=token_embeddings
----

2
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/transformers/TransformersEmbeddingModelAutoConfiguration.java
View File

@@ -52,6 +52,8 @@ public class TransformersEmbeddingModelAutoConfiguration {
embeddingModel.setGpuDeviceId(properties.getOnnx().getGpuDeviceId());
embeddingModel.setModelOutputName(properties.getOnnx().getModelOutputName());
return embeddingModel;
}