From 1746bf6b46c687f0dbe828e6bc1938cb24a2aee5 Mon Sep 17 00:00:00 2001 From: Christian Tzolov Date: Wed, 15 Nov 2023 12:29:22 +0100 Subject: [PATCH] Add transformers-embedding boot auto-configuraiton and starter --- .../transformers-embedding/README.md | 48 ++++- .../transformers-embedding/pom.xml | 8 +- .../TransformersEmbeddingClient.java | 54 ++++- .../TransformersEmbeddingClientTests.java | 2 - spring-ai-huggingface/README.md | 2 +- spring-ai-spring-boot-autoconfigure/pom.xml | 16 +- ...rmersEmbeddingClientAutoConfiguration.java | 57 ++++++ ...TransformersEmbeddingClientProperties.java | 188 ++++++++++++++++++ ...ersEmbeddingClientAutoConfigurationIT.java | 92 +++++++++ .../MilvusVectorStoreAutoConfigurationIT.java | 3 +- .../pom.xml | 51 +++++ .../ai/vectorstore/MilvusVectorStoreIT.java | 5 - 12 files changed, 491 insertions(+), 35 deletions(-) create mode 100644 spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/embedding/transformer/TransformersEmbeddingClientAutoConfiguration.java create mode 100644 spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/embedding/transformer/TransformersEmbeddingClientProperties.java create mode 100644 spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/embedding/transformer/TransformersEmbeddingClientAutoConfigurationIT.java create mode 100644 spring-ai-spring-boot-starters/spring-ai-starter-transformers-embedding/pom.xml diff --git a/embedding-clients/transformers-embedding/README.md b/embedding-clients/transformers-embedding/README.md index 776ca502f..042c662a0 100644 --- a/embedding-clients/transformers-embedding/README.md +++ b/embedding-clients/transformers-embedding/README.md @@ -34,9 +34,9 @@ Add the `transformers-embedding` project to your maven dependencies: ```xml - org.springframework.experimental.ai - transformers-embedding - 0.7.1-SNAPSHOT + org.springframework.experimental.ai + transformers-embedding + 0.7.1-SNAPSHOT ``` @@ -51,8 +51,7 @@ If the model is not explicitly set, `TransformersEmbeddingClient` defaults to [s | Speed | 14200 sentences/sec | | Size | 80MB | - -Following snippet illustrates how to use the `TransformersEmbeddingClient`: +Following snippet illustrates how to use the `TransformersEmbeddingClient` manually: ```java TransformersEmbeddingClient embeddingClient = new TransformersEmbeddingClient(); @@ -73,14 +72,53 @@ List> embeddings = embeddingClient.embed(List.of("Hello world", "Wo ``` +Note that when created manually you have to call the `afterPropertiesSet()` after setting the properties and before using the client. + The first `embed()` call downloads the the large ONNX model and caches it on the local file system. Therefore the first call might take longer than usual. Use the `#setResourceCacheDirectory()` to set the local folder where the ONNX models as stored. The default cache folder is `${java.io.tmpdir}/spring-ai-onnx-model`. +It is more convenient (and preferred) to create the TransformersEmbeddingClient as a `Bean`. +Then you don't have to call the `afterPropertiesSet()` manually. +```java +@Bean +public EmbeddingClient embeddingClient() { + return new TransformersEmbeddingClient(); +} +``` +## Transformers Embedding Spring Boot Starter. +You can bootstrap and auto-wire the `TransformersEmbeddingClient` with following boot starer: +```xml + + org.springframework.experimental.ai + spring-ai-transformers-embedding-spring-boot-starter + 0.7.1-SNAPSHOT + +``` +and use the `spring.ai.embedding.transformer.*` properties to configure it. + +For example add this to your application.properties to configure with the [intfloat/e5-small-v2](https://huggingface.co/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: + +| Property | Description | Default | +| -------- | ------- | ------- | +| 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`' and '`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.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 | diff --git a/embedding-clients/transformers-embedding/pom.xml b/embedding-clients/transformers-embedding/pom.xml index 69f52566e..f79e479a5 100644 --- a/embedding-clients/transformers-embedding/pom.xml +++ b/embedding-clients/transformers-embedding/pom.xml @@ -10,8 +10,8 @@ transformers-embedding jar - Spring AI Embedding Client - Sentence Transormers Embeddings - Spring AI Sentence Transformers Embedding Client + Spring AI Transormers Embedding Client + Spring AI Transformers Embedding Client https://github.com/spring-projects-experimental/spring-ai @@ -21,8 +21,8 @@ - 0.24.0 - 1.16.1 + 0.25.0 + 1.16.2 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 1ae9ad93b..99ad5c78c 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 @@ -5,6 +5,7 @@ import java.util.ArrayList; import java.util.List; import java.util.Map; import java.util.concurrent.atomic.AtomicInteger; +import java.util.stream.Collectors; import ai.djl.huggingface.tokenizers.Encoding; import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer; @@ -17,6 +18,8 @@ import ai.onnxruntime.OnnxValue; import ai.onnxruntime.OrtEnvironment; import ai.onnxruntime.OrtException; import ai.onnxruntime.OrtSession; +import org.apache.commons.logging.Log; +import org.apache.commons.logging.LogFactory; import org.springframework.ai.document.Document; import org.springframework.ai.document.MetadataMode; @@ -33,12 +36,14 @@ import org.springframework.util.StringUtils; */ public class TransformersEmbeddingClient implements EmbeddingClient, InitializingBean { + private static final Log logger = LogFactory.getLog(TransformersEmbeddingClient.class); + // ONNX tokenizer for the all-MiniLM-L6-v2 model - private final static String DEFAULT_ONNX_TOKENIZER_URI = "https://raw.githubusercontent.com/spring-projects-experimental/spring-ai/main/embedding-clients/transformers-embedding/src/main/resources/onnx/all-MiniLM-L6-v2/tokenizer.json"; + public final static String DEFAULT_ONNX_TOKENIZER_URI = "https://raw.githubusercontent.com/spring-projects-experimental/spring-ai/main/embedding-clients/transformers-embedding/src/main/resources/onnx/all-MiniLM-L6-v2/tokenizer.json"; // ONNX model for all-MiniLM-L6-v2 pre-trained transformer: // https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2 - private 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_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"; private final static int EMBEDDING_AXIS = 1; @@ -59,10 +64,19 @@ public class TransformersEmbeddingClient implements EmbeddingClient, Initializin */ private OrtEnvironment environment; + /** + * Runtime session that wraps the ONNX model and enables inference calls. + */ private OrtSession session; private final AtomicInteger embeddingDimensions = new AtomicInteger(-1); + /** + * Specifies what parts of the {@link Document}'s content and metadata will be used + * for computing the embeddings. Applicable for the {@link #embed(Document)} method + * only. Has no effect on the {@link #embed(String)} or {@link #embed(List)}. Defaults + * to {@link MetadataMode#NONE}. + */ private final MetadataMode metadataMode; /** @@ -76,7 +90,15 @@ public class TransformersEmbeddingClient implements EmbeddingClient, Initializin */ private boolean disableCaching = false; - private ResourceCacheService cache; + /** + * Cache service for caching large {@link Resource} contents, such as the + * tokenizerResource and modelResource, on the local file system. Can be + * enabled/disabled with the {@link #disableCaching} property and uses the + * {@link #resourceCacheDirectory} for local storage. + */ + private ResourceCacheService cacheService; + + public Map tokenizerOptions = Map.of(); public TransformersEmbeddingClient() { this(MetadataMode.NONE); @@ -87,6 +109,10 @@ public class TransformersEmbeddingClient implements EmbeddingClient, Initializin this.metadataMode = metadataMode; } + public void setTokenizerOptions(Map tokenizerOptions) { + this.tokenizerOptions = tokenizerOptions; + } + public void setDisableCaching(boolean disableCaching) { this.disableCaching = disableCaching; } @@ -121,23 +147,32 @@ public class TransformersEmbeddingClient implements EmbeddingClient, Initializin @Override public void afterPropertiesSet() throws Exception { - this.cache = StringUtils.hasText(this.resourceCacheDirectory) + + this.cacheService = StringUtils.hasText(this.resourceCacheDirectory) ? new ResourceCacheService(this.resourceCacheDirectory) : new ResourceCacheService(); + + // Create a pre-trained HuggingFaceTokenizer instance from tokenizerResource + // InputStream. this.tokenizer = HuggingFaceTokenizer.newInstance(getCachedResource(this.tokenizerResource).getInputStream(), - Map.of()); + this.tokenizerOptions); + + // onnxruntime this.environment = OrtEnvironment.getEnvironment(); var sessionOptions = new OrtSession.SessionOptions(); if (this.gpuDeviceId >= 0) { - // Run on a GPU or with another provider - sessionOptions.addCUDA(this.gpuDeviceId); + sessionOptions.addCUDA(this.gpuDeviceId); // Run on a GPU or with another + // provider } 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(", "))); } private Resource getCachedResource(Resource resource) { - return this.disableCaching ? resource : this.cache.getCachedResource(resource); + return this.disableCaching ? resource : this.cacheService.getCachedResource(resource); } @Override @@ -186,6 +221,9 @@ public class TransformersEmbeddingClient implements EmbeddingClient, Initializin Map modelInputs = Map.of("input_ids", inputIds, "attention_mask", attentionMask, "token_type_ids", tokenTypeIds); + // 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); diff --git a/embedding-clients/transformers-embedding/src/test/java/org/springframework/ai/embedding/TransformersEmbeddingClientTests.java b/embedding-clients/transformers-embedding/src/test/java/org/springframework/ai/embedding/TransformersEmbeddingClientTests.java index 7fa80f25c..543b1db7d 100644 --- a/embedding-clients/transformers-embedding/src/test/java/org/springframework/ai/embedding/TransformersEmbeddingClientTests.java +++ b/embedding-clients/transformers-embedding/src/test/java/org/springframework/ai/embedding/TransformersEmbeddingClientTests.java @@ -16,11 +16,9 @@ package org.springframework.ai.embedding; -import java.io.File; import java.util.List; import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.io.TempDir; import org.springframework.ai.document.Document; diff --git a/spring-ai-huggingface/README.md b/spring-ai-huggingface/README.md index e4505cf7e..72d43bf10 100644 --- a/spring-ai-huggingface/README.md +++ b/spring-ai-huggingface/README.md @@ -10,7 +10,7 @@ Add the `spring-ai-huggingface` dependency: org.springframework.experimental.ai spring-ai-huggingface - 0.7.0-SNAPSHOT + 0.7.1-SNAPSHOT ``` diff --git a/spring-ai-spring-boot-autoconfigure/pom.xml b/spring-ai-spring-boot-autoconfigure/pom.xml index ac2773f2b..52d128847 100644 --- a/spring-ai-spring-boot-autoconfigure/pom.xml +++ b/spring-ai-spring-boot-autoconfigure/pom.xml @@ -49,6 +49,14 @@ true + + + org.springframework.experimental.ai + transformers-embedding + ${project.parent.version} + true + + org.springframework.experimental.ai @@ -87,7 +95,6 @@ - org.springframework.boot spring-boot-configuration-processor @@ -125,13 +132,6 @@ test - - org.springframework.experimental.ai - transformers-embedding - ${parent.version} - test - - org.testcontainers testcontainers diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/embedding/transformer/TransformersEmbeddingClientAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/embedding/transformer/TransformersEmbeddingClientAutoConfiguration.java new file mode 100644 index 000000000..e646c4d14 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/embedding/transformer/TransformersEmbeddingClientAutoConfiguration.java @@ -0,0 +1,57 @@ +/* + * Copyright 2023-2023 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.ai.autoconfigure.embedding.transformer; + +import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer; +import ai.onnxruntime.OrtSession; + +import org.springframework.ai.embedding.EmbeddingClient; +import org.springframework.ai.embedding.TransformersEmbeddingClient; +import org.springframework.boot.autoconfigure.AutoConfiguration; +import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; +import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; +import org.springframework.boot.context.properties.EnableConfigurationProperties; +import org.springframework.context.annotation.Bean; + +/** + * @author Christian Tzolov + */ +@AutoConfiguration +@EnableConfigurationProperties({ TransformersEmbeddingClientProperties.class }) +@ConditionalOnClass({ OrtSession.class, HuggingFaceTokenizer.class }) +public class TransformersEmbeddingClientAutoConfiguration { + + @Bean + @ConditionalOnMissingBean + public EmbeddingClient embeddingClient(TransformersEmbeddingClientProperties properties) { + + TransformersEmbeddingClient embeddingClient = new TransformersEmbeddingClient(properties.getMetadataMode()); + + embeddingClient.setDisableCaching(!properties.getCache().isEnabled()); + embeddingClient.setResourceCacheDirectory(properties.getCache().getDirectory()); + + embeddingClient.setTokenizerResource(properties.getTokenizer().getUri()); + embeddingClient.setTokenizerOptions(properties.getTokenizer().getOptions()); + + embeddingClient.setModelResource(properties.getOnnx().getModelUri()); + + embeddingClient.setGpuDeviceId(properties.getOnnx().getGpuDeviceId()); + + return embeddingClient; + } + +} 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 new file mode 100644 index 000000000..d58efc0b1 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/embedding/transformer/TransformersEmbeddingClientProperties.java @@ -0,0 +1,188 @@ +/* + * Copyright 2023-2023 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.ai.autoconfigure.embedding.transformer; + +import java.io.File; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer; + +import org.springframework.ai.document.Document; +import org.springframework.ai.document.MetadataMode; +import org.springframework.ai.embedding.TransformersEmbeddingClient; +import org.springframework.boot.context.properties.ConfigurationProperties; + +import static org.springframework.ai.autoconfigure.embedding.transformer.TransformersEmbeddingClientProperties.CONFIG_PREFIX; + +/** + * @author Christian Tzolov + */ +@ConfigurationProperties(CONFIG_PREFIX) +public class TransformersEmbeddingClientProperties { + + public static final String CONFIG_PREFIX = "spring.ai.embedding.transformer"; + + public static final String DEFAULT_CACHE_DIRECTORY = new File(System.getProperty("java.io.tmpdir"), + "spring-ai-onnx-model") + .getAbsolutePath(); + + /** + * Configurations for the {@link HuggingFaceTokenizer} used to convert sentences into + * tokens. + */ + public static class Tokenizer { + + /** + * URI of a pre-trained HuggingFaceTokenizer created by the ONNX engine (e.g. + * tokenizer.json). + */ + private String uri = TransformersEmbeddingClient.DEFAULT_ONNX_TOKENIZER_URI; + + /** + * HuggingFaceTokenizer options such as 'addSpecialTokens', 'modelMaxLength', + * 'truncation', 'padding', 'maxLength', 'stride' and 'padToMultipleOf'. Leave + * empty to fallback to the defaults. + */ + private Map options = new HashMap<>(); + + public String getUri() { + return uri; + } + + public void setUri(String uri) { + this.uri = uri; + } + + public Map getOptions() { + return options; + } + + public void setOptions(Map options) { + this.options = options; + } + + } + + private final Tokenizer tokenizer = new Tokenizer(); + + public static class Cache { + + /** + * Enable the {@link Resource} caching. + */ + private boolean enabled = true; + + /** + * Resource cache directory. Used to cache remote resources, such as the ONNX + * models, to the local file system. Applicable only for cache.enabled == true. + * Defaults to {java.io.tmpdir}/spring-ai-onnx-model. + */ + private String directory = DEFAULT_CACHE_DIRECTORY; + + public boolean isEnabled() { + return enabled; + } + + public void setEnabled(boolean enabled) { + this.enabled = enabled; + } + + public String getDirectory() { + return directory; + } + + public void setDirectory(String directory) { + this.directory = directory; + } + + } + + /** + * Controls caching of remote, large resources to local file system. + */ + private final Cache cache = new Cache(); + + public Cache getCache() { + return cache; + } + + public static class Onnx { + + /** + * Existing, pre-trained ONNX model. Commonly exported from + * https://sbert.net/docs/pretrained_models.html. Defaults to + * sentence-transformers/all-MiniLM-L6-v2. + */ + private String modelUri = TransformersEmbeddingClient.DEFAULT_ONNX_MODEL_URI; + + /** + * 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 + * + * The GPU device ID to execute on. Only applicable if >= 0. Ignored otherwise. + */ + private int gpuDeviceId = -1; + + public String getModelUri() { + return modelUri; + } + + public void setModelUri(String modelUri) { + this.modelUri = modelUri; + } + + public int getGpuDeviceId() { + return gpuDeviceId; + } + + public void setGpuDeviceId(int gpuDeviceId) { + this.gpuDeviceId = gpuDeviceId; + } + + } + + private final Onnx onnx = new Onnx(); + + public Onnx getOnnx() { + return onnx; + } + + /** + * Specifies what parts of the {@link Document}'s content and metadata will be used + * for computing the embeddings. Applicable for the + * {@link TransformersEmbeddingClient#embed(Document)} method only. Has no effect on + * the {@link TransformersEmbeddingClient#embed(String)} or + * {@link TransformersEmbeddingClient#embed(List)}. Defaults to + * {@link MetadataMode#NONE}. + */ + private MetadataMode metadataMode = MetadataMode.NONE; + + public Tokenizer getTokenizer() { + return tokenizer; + } + + public MetadataMode getMetadataMode() { + return metadataMode; + } + + public void setMetadataMode(MetadataMode metadataMode) { + this.metadataMode = metadataMode; + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/embedding/transformer/TransformersEmbeddingClientAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/embedding/transformer/TransformersEmbeddingClientAutoConfigurationIT.java new file mode 100644 index 000000000..84ba40572 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/embedding/transformer/TransformersEmbeddingClientAutoConfigurationIT.java @@ -0,0 +1,92 @@ +/* + * Copyright 2023-2023 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.ai.autoconfigure.embedding.transformer; + +import java.io.File; +import java.util.List; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import org.springframework.ai.embedding.EmbeddingClient; +import org.springframework.ai.embedding.TransformersEmbeddingClient; +import org.springframework.boot.autoconfigure.AutoConfigurations; +import org.springframework.boot.test.context.runner.ApplicationContextRunner; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Christian Tzolov + */ +public class TransformersEmbeddingClientAutoConfigurationIT { + + @TempDir + File tempDir; + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withConfiguration(AutoConfigurations.of(TransformersEmbeddingClientAutoConfiguration.class)); + + @Test + public void embedding() { + contextRunner.run(context -> { + var properties = context.getBean(TransformersEmbeddingClientProperties.class); + assertThat(properties.getCache().isEnabled()).isTrue(); + assertThat(properties.getCache().getDirectory()) + .isEqualTo(new File(System.getProperty("java.io.tmpdir"), "spring-ai-onnx-model").getAbsolutePath()); + + EmbeddingClient embeddingClient = context.getBean(EmbeddingClient.class); + assertThat(embeddingClient).isInstanceOf(TransformersEmbeddingClient.class); + + List> embeddings = embeddingClient.embed(List.of("Spring Framework", "Spring AI")); + + assertThat(embeddings.size()).isEqualTo(2); // batch size + assertThat(embeddings.get(0).size()).isEqualTo(embeddingClient.dimensions()); // dimensions + // size + }); + } + + @Test + public void remoteOnnxModel() { + // https://huggingface.co/intfloat/e5-small-v2 + contextRunner.withPropertyValues("spring.ai.embedding.transformer.cache.directory=" + tempDir.getAbsolutePath(), + "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") + .run(context -> { + var properties = context.getBean(TransformersEmbeddingClientProperties.class); + assertThat(properties.getOnnx().getModelUri()) + .isEqualTo("https://huggingface.co/intfloat/e5-small-v2/resolve/main/model.onnx"); + assertThat(properties.getTokenizer().getUri()) + .isEqualTo("https://huggingface.co/intfloat/e5-small-v2/raw/main/tokenizer.json"); + + assertThat(properties.getCache().isEnabled()).isTrue(); + assertThat(properties.getCache().getDirectory()).isEqualTo(tempDir.getAbsolutePath()); + assertThat(tempDir.listFiles()).hasSize(2); + + EmbeddingClient embeddingClient = context.getBean(EmbeddingClient.class); + assertThat(embeddingClient).isInstanceOf(TransformersEmbeddingClient.class); + + assertThat(embeddingClient.dimensions()).isEqualTo(384); + + List> embeddings = embeddingClient.embed(List.of("Spring Framework", "Spring AI")); + + assertThat(embeddings.size()).isEqualTo(2); // batch size + assertThat(embeddings.get(0).size()).isEqualTo(embeddingClient.dimensions()); // dimensions + // size + }); + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/milvus/MilvusVectorStoreAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/milvus/MilvusVectorStoreAutoConfigurationIT.java index 313e6b7f6..7bec6324f 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/milvus/MilvusVectorStoreAutoConfigurationIT.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/milvus/MilvusVectorStoreAutoConfigurationIT.java @@ -27,12 +27,10 @@ import java.util.UUID; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; -import org.springframework.core.io.DefaultResourceLoader; import org.testcontainers.containers.DockerComposeContainer; import org.testcontainers.containers.wait.strategy.Wait; import org.testcontainers.junit.jupiter.Testcontainers; -import org.springframework.ai.ResourceUtils; import org.springframework.ai.document.Document; import org.springframework.ai.embedding.EmbeddingClient; import org.springframework.ai.embedding.TransformersEmbeddingClient; @@ -42,6 +40,7 @@ import org.springframework.boot.autoconfigure.AutoConfigurations; import org.springframework.boot.test.context.runner.ApplicationContextRunner; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; +import org.springframework.core.io.DefaultResourceLoader; import org.springframework.util.FileSystemUtils; import static org.assertj.core.api.Assertions.assertThat; diff --git a/spring-ai-spring-boot-starters/spring-ai-starter-transformers-embedding/pom.xml b/spring-ai-spring-boot-starters/spring-ai-starter-transformers-embedding/pom.xml new file mode 100644 index 000000000..c4af13a54 --- /dev/null +++ b/spring-ai-spring-boot-starters/spring-ai-starter-transformers-embedding/pom.xml @@ -0,0 +1,51 @@ + + + 4.0.0 + + org.springframework.experimental.ai + spring-ai + 0.7.1-SNAPSHOT + ../../pom.xml + + spring-ai-transformers-embedding-spring-boot-starter + jar + Spring AI Starter - Transformers Embedding + Spring Transformers Embedding Auto Configuration + https://github.com/spring-projects-experimental/spring-ai + + + https://github.com/spring-projects-experimental/spring-ai + git://github.com/spring-projects-experimental/spring-ai.git + git@github.com:spring-projects-experimental/spring-ai.git + + + + + + + + org.springframework.boot + spring-boot-starter + + + + org.springframework.experimental.ai + spring-ai-spring-boot-autoconfigure + ${project.parent.version} + + + + org.springframework.experimental.ai + transformers-embedding + ${project.parent.version} + + + + + org.springframework.boot + spring-boot-starter-test + test + + + + diff --git a/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreIT.java b/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreIT.java index 5081ca53b..139c76de1 100644 --- a/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreIT.java +++ b/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreIT.java @@ -310,11 +310,6 @@ public class MilvusVectorStoreIT { return new OpenAiEmbeddingClient(new OpenAiService(api), "text-embedding-ada-002"); } - // @Bean - // public EmbeddingClient embeddingClient() { - // return new TransformersEmbeddingClient(); - // } - } } \ No newline at end of file