Transformers-embedding client - align tokenizer with onnx's model inputs

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<List<Double>> 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
+```

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<String, String> tokenizerOptions = Map.of();
 
+	private String modelOutputName = DEFAULT_MODEL_OUTPUT_NAME;
+
+	private Set<String> 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<String> 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<String, OnnxTensor> 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<String, OnnxTensor> removeUnknownModelInputs(Map<String, OnnxTensor> 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) {
 

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