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

Ollama and Titan Embedding Client makes individual calls to embed Document collection

Browse Source 
Fixes #218
This commit is contained in:
Mark Pollack
2024-01-11 12:15:52 -05:00
parent 85e3e0015e
commit 408d75e0d4
4 changed files with 34 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingClient.java
View File

@@ -16,9 +16,12 @@
package org.springframework.ai.bedrock.titan;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi;
import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi.TitanEmbeddingRequest;
import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi.TitanEmbeddingResponse;
@@ -41,6 +44,8 @@ import org.springframework.util.Assert;
*/
public class BedrockTitanEmbeddingClient extends AbstractEmbeddingClient {
private final Logger logger = LoggerFactory.getLogger(getClass());
private final TitanEmbeddingBedrockApi embeddingApi;
public enum InputType {
@@ -90,17 +95,20 @@ public class BedrockTitanEmbeddingClient extends AbstractEmbeddingClient {
@Override
public List<List<Double>> embed(List<String> inputContents) {
Assert.notEmpty(inputContents, "At least one text is required!");
Assert.isTrue(inputContents.size() == 1, "Titan Embedding does not support batch embedding!");
if (inputContents.size() != 1) {
logger.warn(
"Titan Embedding does not support batch embedding. Will make multiple API calls to embed(Document)");
}
String inputContent = inputContents.iterator().next();
var request = (this.inputType == InputType.IMAGE)
? new TitanEmbeddingRequest.Builder().withInputImage(inputContent).build()
: new TitanEmbeddingRequest.Builder().withInputText(inputContent).build();
TitanEmbeddingResponse response = this.embeddingApi.embedding(request);
return List.of(response.embedding());
List<List<Double>> embeddingList = new ArrayList<>();
for (String inputContent : inputContents) {
var request = (this.inputType == InputType.IMAGE)
? new TitanEmbeddingRequest.Builder().withInputImage(inputContent).build()
: new TitanEmbeddingRequest.Builder().withInputText(inputContent).build();
TitanEmbeddingResponse response = this.embeddingApi.embedding(request);
embeddingList.add(response.embedding());
}
return embeddingList;
}
@Override

8
models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingClientIT.java
View File

@@ -37,14 +37,6 @@ class BedrockTitanEmbeddingClientIT {
assertThat(embeddingClient.dimensions()).isEqualTo(1024);
}
@Test
void batchEmbedding() {
assertThatThrownBy(
() -> embeddingClient.embedForResponse(List.of("Hello World", "World is big and salvation is near")))
.isInstanceOf(IllegalArgumentException.class)
.hasMessageContaining("Titan Embedding does not support batch embedding!");
}
@Test
void imageEmbedding() throws IOException {

23
models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaEmbeddingClient.java
View File

@@ -16,10 +16,13 @@
package org.springframework.ai.ollama;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.AbstractEmbeddingClient;
import org.springframework.ai.embedding.Embedding;
@@ -49,6 +52,8 @@ import org.springframework.util.Assert;
*/
public class OllamaEmbeddingClient extends AbstractEmbeddingClient {
private final Logger logger = LoggerFactory.getLogger(getClass());
private final OllamaApi ollamaApi;
private String model = "orca-mini";
@@ -87,14 +92,18 @@ public class OllamaEmbeddingClient extends AbstractEmbeddingClient {
@Override
public List<List<Double>> embed(List<String> texts) {
Assert.notEmpty(texts, "At least one text is required!");
Assert.isTrue(texts.size() == 1, "Ollama Embedding does not support batch embedding!");
if (texts.size() != 1) {
logger.warn(
"Ollama Embedding does not support batch embedding. Will make multiple API calls to embed(Document)");
}
String inputContent = texts.iterator().next();
OllamaApi.EmbeddingResponse response = this.ollamaApi
.embeddings(new EmbeddingRequest(this.model, inputContent, this.clientOptions));
return List.of(response.embedding());
List<List<Double>> embeddingList = new ArrayList<>();
for (String inputContent : texts) {
OllamaApi.EmbeddingResponse response = this.ollamaApi
.embeddings(new EmbeddingRequest(this.model, inputContent, this.clientOptions));
embeddingList.add(response.embedding());
}
return embeddingList;
}
@Override

8
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaEmbeddingClientIT.java
View File

@@ -56,14 +56,6 @@ class OllamaEmbeddingClientIT {
assertThat(embeddingClient.dimensions()).isEqualTo(3200);
}
@Test
void batchEmbedding() {
assertThatThrownBy(
() -> embeddingClient.embedForResponse(List.of("Hello World", "World is big and salvation is near")))
.isInstanceOf(IllegalArgumentException.class)
.hasMessageContaining("Ollama Embedding does not support batch embedding!");
}
@SpringBootConfiguration
public static class TestConfiguration {