From 3ea938c4e12687efe8ac998b7ce204a90792e81c Mon Sep 17 00:00:00 2001 From: Christian Tzolov Date: Wed, 15 Nov 2023 15:18:22 +0100 Subject: [PATCH] Transformers-embedding client - align tokenizer with onnx's model inputs --- .../transformers-embedding/README.md | 10 ++++++ .../TransformersEmbeddingClient.java | 34 +++++++++++++++++-- ...TransformersEmbeddingClientProperties.java | 15 +++++++- 3 files changed, 55 insertions(+), 4 deletions(-) diff --git a/embedding-clients/transformers-embedding/README.md b/embedding-clients/transformers-embedding/README.md index 042c662a0..0ce2eeb13 100644 --- a/embedding-clients/transformers-embedding/README.md +++ b/embedding-clients/transformers-embedding/README.md @@ -66,6 +66,10 @@ embeddingClient.setModelResource("classpath:/onnx/all-MiniLM-L6-v2/model.onnx"); // 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> embeddings = embeddingClient.embed(List.of("Hello world", "World is big")); @@ -122,3 +126,9 @@ The complete list of supported properties are: | 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 error like: `Caused by: ai.onnxruntime.OrtException: Supplied array is ragged,..` then you need to enable the tokenizer padding in boot starter's `application.properties`: + +``` +spring.ai.embedding.transformer.tokenizer.options.padding=true +``` \ No newline at end of file diff --git a/embedding-clients/transformers-embedding/src/main/java/org/springframework/ai/embedding/TransformersEmbeddingClient.java b/embedding-clients/transformers-embedding/src/main/java/org/springframework/ai/embedding/TransformersEmbeddingClient.java index 99ad5c78c..87e52858c 100644 --- a/embedding-clients/transformers-embedding/src/main/java/org/springframework/ai/embedding/TransformersEmbeddingClient.java +++ b/embedding-clients/transformers-embedding/src/main/java/org/springframework/ai/embedding/TransformersEmbeddingClient.java @@ -4,6 +4,7 @@ import java.nio.FloatBuffer; import java.util.ArrayList; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; @@ -45,6 +46,8 @@ public class TransformersEmbeddingClient implements EmbeddingClient, Initializin // https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2 public final static String DEFAULT_ONNX_MODEL_URI = "https://github.com/spring-projects-experimental/spring-ai/raw/main/embedding-clients/transformers-embedding/src/main/resources/onnx/all-MiniLM-L6-v2/model.onnx"; + public final static String DEFAULT_MODEL_OUTPUT_NAME = "last_hidden_state"; + private final static int EMBEDDING_AXIS = 1; private Resource tokenizerResource = toResource(DEFAULT_ONNX_TOKENIZER_URI); @@ -100,6 +103,10 @@ public class TransformersEmbeddingClient implements EmbeddingClient, Initializin public Map tokenizerOptions = Map.of(); + private String modelOutputName = DEFAULT_MODEL_OUTPUT_NAME; + + private Set onnxModelInputs; + public TransformersEmbeddingClient() { this(MetadataMode.NONE); } @@ -145,6 +152,10 @@ public class TransformersEmbeddingClient implements EmbeddingClient, Initializin this.embeddingDimensions.set(dimension); } + public void setModelOutputName(String modelOutputName) { + this.modelOutputName = modelOutputName; + } + @Override public void afterPropertiesSet() throws Exception { @@ -167,8 +178,14 @@ public class TransformersEmbeddingClient implements EmbeddingClient, Initializin this.session = this.environment.createSession(getCachedResource(this.modelResource).getContentAsByteArray(), sessionOptions); - logger.info("Model input names: " + this.session.getInputNames().stream().collect(Collectors.joining(", "))); - logger.info("Model output names: " + this.session.getOutputNames().stream().collect(Collectors.joining(", "))); + this.onnxModelInputs = this.session.getInputNames(); + Set onnxModelOutputs = this.session.getOutputNames(); + + logger.info("Model input names: " + this.onnxModelInputs.stream().collect(Collectors.joining(", "))); + logger.info("Model output names: " + onnxModelOutputs.stream().collect(Collectors.joining(", "))); + + Assert.isTrue(onnxModelOutputs.contains(this.modelOutputName), + "The model output names doesn't contain expected: " + this.modelOutputName); } private Resource getCachedResource(Resource resource) { @@ -221,13 +238,15 @@ public class TransformersEmbeddingClient implements EmbeddingClient, Initializin Map modelInputs = Map.of("input_ids", inputIds, "attention_mask", attentionMask, "token_type_ids", tokenTypeIds); + modelInputs = removeUnknownModelInputs(modelInputs); + // The Run result object is AutoCloseable to prevent references from leaking // out. Once the Result object is // closed, all it’s child OnnxValues are closed too. try (OrtSession.Result results = this.session.run(modelInputs)) { // OnnxValue lastHiddenState = results.get(0); - OnnxValue lastHiddenState = results.get("last_hidden_state").get(); + OnnxValue lastHiddenState = results.get(this.modelOutputName).get(); // 0 - batch_size (1..x) // 1 - sequence_length (128) @@ -253,6 +272,15 @@ public class TransformersEmbeddingClient implements EmbeddingClient, Initializin return resultEmbeddings; } + private Map removeUnknownModelInputs(Map modelInputs) { + + return modelInputs.entrySet() + .stream() + .filter(a -> onnxModelInputs.contains(a.getKey())) + .collect(Collectors.toMap(e -> e.getKey(), e -> e.getValue())); + + } + // Build a NDArray from 3D float array. private NDArray create(float[][][] data3d, NDManager manager) { diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/embedding/transformer/TransformersEmbeddingClientProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/embedding/transformer/TransformersEmbeddingClientProperties.java index d58efc0b1..c81e7dd18 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/embedding/transformer/TransformersEmbeddingClientProperties.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/embedding/transformer/TransformersEmbeddingClientProperties.java @@ -84,7 +84,7 @@ public class TransformersEmbeddingClientProperties { public static class Cache { /** - * Enable the {@link Resource} caching. + * Enable the Resource caching. */ private boolean enabled = true; @@ -131,6 +131,11 @@ public class TransformersEmbeddingClientProperties { */ private String modelUri = TransformersEmbeddingClient.DEFAULT_ONNX_MODEL_URI; + /** + * Defaults to: 'last_hidden_state'. + */ + private String modelOutputName = TransformersEmbeddingClient.DEFAULT_MODEL_OUTPUT_NAME; + /** * Run on a GPU or with another provider (optional). * https://onnxruntime.ai/docs/get-started/with-java.html#run-on-a-gpu-or-with-another-provider-optional @@ -155,6 +160,14 @@ public class TransformersEmbeddingClientProperties { this.gpuDeviceId = gpuDeviceId; } + public String getModelOutputName() { + return modelOutputName; + } + + public void setModelOutputName(String modelOutputName) { + this.modelOutputName = modelOutputName; + } + } private final Onnx onnx = new Onnx();