Refactor Ollama embedding model implementation

- Update OllamaEmbeddingModel to support batch embedding requests
- Rename EmbeddingRequest/Response to EmbeddingsRequest/Response
- Add truncate option to control input truncation
- Update documentation and tests for new embedding API
- Remove deprecated withModel and withDefaultOptions methods
- Adjust default values for various Ollama model options
- Include response medata with response model name

Resolves #1158
This commit is contained in:
Christian Tzolov
2024-08-03 22:06:55 +02:00
parent 3978e8ecc2
commit 4cacbe8468
10 changed files with 372 additions and 157 deletions

View File

@@ -15,22 +15,26 @@
*/
package org.springframework.ai.ollama;
import java.util.ArrayList;
import java.time.Duration;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.metadata.EmptyUsage;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.AbstractEmbeddingModel;
import org.springframework.ai.embedding.Embedding;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.EmbeddingOptions;
import org.springframework.ai.embedding.EmbeddingRequest;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.embedding.EmbeddingResponseMetadata;
import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaApi.EmbeddingRequest;
import org.springframework.ai.ollama.api.OllamaApi.EmbeddingsResponse;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;
@@ -71,70 +75,43 @@ public class OllamaEmbeddingModel extends AbstractEmbeddingModel {
this.defaultOptions = defaultOptions;
}
/**
* @deprecated Use {@link OllamaOptions#setModel} instead.
*/
@Deprecated
public OllamaEmbeddingModel withModel(String model) {
this.defaultOptions.setModel(model);
return this;
}
/**
* @deprecated Use {@link OllamaOptions} constructor instead.
*/
@Deprecated
public OllamaEmbeddingModel withDefaultOptions(OllamaOptions options) {
this.defaultOptions = options;
return this;
}
@Override
public List<Double> embed(Document document) {
return embed(document.getContent());
}
@Override
public EmbeddingResponse call(org.springframework.ai.embedding.EmbeddingRequest request) {
public EmbeddingResponse call(EmbeddingRequest request) {
Assert.notEmpty(request.getInstructions(), "At least one text is required!");
if (request.getInstructions().size() != 1) {
logger.warn(
"Ollama Embedding does not support batch embedding. Will make multiple API calls to embed(Document)");
}
List<List<Double>> embeddingList = new ArrayList<>();
for (String inputContent : request.getInstructions()) {
OllamaApi.EmbeddingsRequest ollamaEmbeddingRequest = ollamaEmbeddingRequest(request.getInstructions(),
request.getOptions());
EmbeddingRequest ollamaEmbeddingRequest = ollamaEmbeddingRequest(inputContent, request.getOptions());
EmbeddingsResponse response = this.ollamaApi.embed(ollamaEmbeddingRequest);
OllamaApi.EmbeddingResponse response = this.ollamaApi.embeddings(ollamaEmbeddingRequest);
embeddingList.add(response.embedding());
}
AtomicInteger indexCounter = new AtomicInteger(0);
List<Embedding> embeddings = embeddingList.stream()
List<Embedding> embeddings = response.embeddings()
.stream()
.map(e -> new Embedding(e, indexCounter.getAndIncrement()))
.toList();
return new EmbeddingResponse(embeddings);
EmbeddingResponseMetadata embeddingResponseMetadata = new EmbeddingResponseMetadata(response.model(),
new EmptyUsage());
return new EmbeddingResponse(embeddings, embeddingResponseMetadata);
}
/**
* Package access for testing.
*/
OllamaApi.EmbeddingRequest ollamaEmbeddingRequest(String inputContent, EmbeddingOptions options) {
OllamaApi.EmbeddingsRequest ollamaEmbeddingRequest(List<String> inputContent, EmbeddingOptions options) {
// runtime options
OllamaOptions runtimeOptions = null;
if (options != null) {
if (options instanceof OllamaOptions ollamaOptions) {
runtimeOptions = ollamaOptions;
}
else {
// currently EmbeddingOptions does not have any portable options to be
// merged.
runtimeOptions = null;
}
if (options != null && options instanceof OllamaOptions ollamaOptions) {
runtimeOptions = ollamaOptions;
}
OllamaOptions mergedOptions = ModelOptionsUtils.merge(runtimeOptions, this.defaultOptions, OllamaOptions.class);
@@ -144,8 +121,40 @@ public class OllamaEmbeddingModel extends AbstractEmbeddingModel {
throw new IllegalArgumentException("Model is not set!");
}
String model = mergedOptions.getModel();
return new EmbeddingRequest(model, inputContent, null,
OllamaOptions.filterNonSupportedFields(mergedOptions.toMap()));
return new OllamaApi.EmbeddingsRequest(model, inputContent, DurationParser.parse(mergedOptions.getKeepAlive()),
OllamaOptions.filterNonSupportedFields(mergedOptions.toMap()), mergedOptions.getTruncate());
}
public static class DurationParser {
private static Pattern PATTERN = Pattern.compile("(\\d+)(ms|s|m|h)");
public static Duration parse(String input) {
if (!StringUtils.hasText(input)) {
return null;
}
Matcher matcher = PATTERN.matcher(input);
if (matcher.matches()) {
long value = Long.parseLong(matcher.group(1));
String unit = matcher.group(2);
return switch (unit) {
case "ms" -> Duration.ofMillis(value);
case "s" -> Duration.ofSeconds(value);
case "m" -> Duration.ofMinutes(value);
case "h" -> Duration.ofHours(value);
default -> throw new IllegalArgumentException("Unsupported time unit: " + unit);
};
}
else {
throw new IllegalArgumentException("Invalid duration format: " + input);
}
}
}
}

View File

@@ -307,7 +307,7 @@ public class OllamaApi {
@JsonProperty("eval_duration") Duration evalDuration) {
}
/**
/**
* Generate a completion for the given prompt.
* @param completionRequest Completion request.
* @return Completion response.
@@ -691,12 +691,41 @@ public class OllamaApi {
* Generate embeddings from a model.
*
* @param model The name of model to generate embeddings from.
* @param prompt The text to generate embeddings for.
* @param input The text or list of text to generate embeddings for.
* @param keepAlive Controls how long the model will stay loaded into memory following the request (default: 5m).
* @param options Additional model parameters listed in the documentation for the
* Model file such as temperature.
* @param truncate Truncates the end of each input to fit within context length.
* Returns error if false and context length is exceeded. Defaults to true.
*/
@JsonInclude(Include.NON_NULL)
public record EmbeddingsRequest(
@JsonProperty("model") String model,
@JsonProperty("input") List<String> input,
@JsonProperty("keep_alive") Duration keepAlive,
@JsonProperty("options") Map<String, Object> options,
@JsonProperty("truncate") Boolean truncate) {
/**
* Shortcut constructor to create a EmbeddingRequest without options.
* @param model The name of model to generate embeddings from.
* @param input The text or list of text to generate embeddings for.
*/
public EmbeddingsRequest(String model, String input) {
this(model, List.of(input), null, null, null);
}
}
/**
* Generate embeddings from a model.
*
* @param model The name of model to generate embeddings from.
* @param prompt The text generate embeddings for
* @param keepAlive Controls how long the model will stay loaded into memory following the request (default: 5m).
* @param options Additional model parameters listed in the documentation for the
* @deprecated Use {@link EmbeddingsRequest} instead.
*/
@Deprecated(since = "1.0.0-M2", forRemoval = true)
@JsonInclude(Include.NON_NULL)
public record EmbeddingRequest(
@JsonProperty("model") String model,
@JsonProperty("prompt") String prompt,
@@ -717,17 +746,49 @@ public class OllamaApi {
* The response object returned from the /embedding endpoint.
*
* @param embedding The embedding generated from the model.
* @deprecated Use {@link EmbeddingsResponse} instead.
*/
@Deprecated(since = "1.0.0-M2", forRemoval = true)
@JsonInclude(Include.NON_NULL)
public record EmbeddingResponse(
@JsonProperty("embedding") List<Double> embedding) {
}
/**
* The response object returned from the /embedding endpoint.
* @param model The model used for generating the embeddings.
* @param embeddings The list of embeddings generated from the model.
* Each embedding (list of doubles) corresponds to a single input text.
*/
@JsonInclude(Include.NON_NULL)
public record EmbeddingsResponse(
@JsonProperty("model") String model,
@JsonProperty("embeddings") List<List<Double>> embeddings) {
}
/**
* Generate embeddings from a model.
* @param embeddingsRequest Embedding request.
* @return Embeddings response.
*/
public EmbeddingsResponse embed(EmbeddingsRequest embeddingsRequest) {
Assert.notNull(embeddingsRequest, REQUEST_BODY_NULL_ERROR);
return this.restClient.post()
.uri("/api/embed")
.body(embeddingsRequest)
.retrieve()
.onStatus(this.responseErrorHandler)
.body(EmbeddingsResponse.class);
}
/**
* Generate embeddings from a model.
* @param embeddingRequest Embedding request.
* @return Embedding response.
* @deprecated Use {@link #embed(EmbeddingsRequest)} instead.
*/
@Deprecated(since = "1.0.0-M2", forRemoval = true)
public EmbeddingResponse embeddings(EmbeddingRequest embeddingRequest) {
Assert.notNull(embeddingRequest, REQUEST_BODY_NULL_ERROR);

View File

@@ -52,7 +52,7 @@ public class OllamaOptions implements FunctionCallingOptions, ChatOptions, Embed
public static final String DEFAULT_MODEL = OllamaModel.MISTRAL.id();
private static final List<String> NON_SUPPORTED_FIELDS = List.of("model", "format", "keep_alive");
private static final List<String> NON_SUPPORTED_FIELDS = List.of("model", "format", "keep_alive", "truncate");
// Following fields are options which must be set when the model is loaded into
// memory.
@@ -267,6 +267,13 @@ public class OllamaOptions implements FunctionCallingOptions, ChatOptions, Embed
* Part of Chat completion <a href="https://github.com/ollama/ollama/blob/main/docs/api.md#parameters-1">advanced parameters</a>.
*/
@JsonProperty("keep_alive") private String keepAlive;
/**
* Truncates the end of each input to fit within context length. Returns error if false and context length is exceeded.
* Defaults to true.
*/
@JsonProperty("truncate") private Boolean truncate;
/**
* Tool Function Callbacks to register with the ChatModel.
@@ -312,14 +319,6 @@ public class OllamaOptions implements FunctionCallingOptions, ChatOptions, Embed
return this;
}
public String getModel() {
return model;
}
public void setModel(String model) {
this.model = model;
}
public OllamaOptions withFormat(String format) {
this.format = format;
return this;
@@ -330,6 +329,11 @@ public class OllamaOptions implements FunctionCallingOptions, ChatOptions, Embed
return this;
}
public OllamaOptions withTruncate(Boolean truncate) {
this.truncate = truncate;
return this;
}
public OllamaOptions withUseNUMA(Boolean useNUMA) {
this.useNUMA = useNUMA;
return this;
@@ -491,6 +495,17 @@ public class OllamaOptions implements FunctionCallingOptions, ChatOptions, Embed
return this;
}
// -------------------
// Getters and Setters
// -------------------
public String getModel() {
return model;
}
public void setModel(String model) {
this.model = model;
}
public String getFormat() {
return this.format;
}
@@ -739,6 +754,14 @@ public class OllamaOptions implements FunctionCallingOptions, ChatOptions, Embed
this.stop = stop;
}
public Boolean getTruncate() {
return this.truncate;
}
public void setTruncate(Boolean truncate) {
this.truncate = truncate;
}
@Override
public List<FunctionCallback> getFunctionCallbacks() {
return this.functionCallbacks;
@@ -797,6 +820,7 @@ public class OllamaOptions implements FunctionCallingOptions, ChatOptions, Embed
.withModel(fromOptions.getModel())
.withFormat(fromOptions.getFormat())
.withKeepAlive(fromOptions.getKeepAlive())
.withTruncate(fromOptions.getTruncate())
.withUseNUMA(fromOptions.getUseNUMA())
.withNumCtx(fromOptions.getNumCtx())
.withNumBatch(fromOptions.getNumBatch())
@@ -839,15 +863,16 @@ public class OllamaOptions implements FunctionCallingOptions, ChatOptions, Embed
return false;
OllamaOptions that = (OllamaOptions) o;
return Objects.equals(model, that.model) && Objects.equals(format, that.format)
&& Objects.equals(keepAlive, that.keepAlive) && Objects.equals(useNUMA, that.useNUMA)
&& Objects.equals(numCtx, that.numCtx) && Objects.equals(numBatch, that.numBatch)
&& Objects.equals(numGPU, that.numGPU) && Objects.equals(mainGPU, that.mainGPU)
&& Objects.equals(lowVRAM, that.lowVRAM) && Objects.equals(f16KV, that.f16KV)
&& Objects.equals(logitsAll, that.logitsAll) && Objects.equals(vocabOnly, that.vocabOnly)
&& Objects.equals(useMMap, that.useMMap) && Objects.equals(useMLock, that.useMLock)
&& Objects.equals(numThread, that.numThread) && Objects.equals(numKeep, that.numKeep)
&& Objects.equals(seed, that.seed) && Objects.equals(numPredict, that.numPredict)
&& Objects.equals(topK, that.topK) && Objects.equals(topP, that.topP) && Objects.equals(tfsZ, that.tfsZ)
&& Objects.equals(keepAlive, that.keepAlive) && Objects.equals(truncate, that.truncate)
&& Objects.equals(useNUMA, that.useNUMA) && Objects.equals(numCtx, that.numCtx)
&& Objects.equals(numBatch, that.numBatch) && Objects.equals(numGPU, that.numGPU)
&& Objects.equals(mainGPU, that.mainGPU) && Objects.equals(lowVRAM, that.lowVRAM)
&& Objects.equals(f16KV, that.f16KV) && Objects.equals(logitsAll, that.logitsAll)
&& Objects.equals(vocabOnly, that.vocabOnly) && Objects.equals(useMMap, that.useMMap)
&& Objects.equals(useMLock, that.useMLock) && Objects.equals(numThread, that.numThread)
&& Objects.equals(numKeep, that.numKeep) && Objects.equals(seed, that.seed)
&& Objects.equals(numPredict, that.numPredict) && Objects.equals(topK, that.topK)
&& Objects.equals(topP, that.topP) && Objects.equals(tfsZ, that.tfsZ)
&& Objects.equals(typicalP, that.typicalP) && Objects.equals(repeatLastN, that.repeatLastN)
&& Objects.equals(temperature, that.temperature) && Objects.equals(repeatPenalty, that.repeatPenalty)
&& Objects.equals(presencePenalty, that.presencePenalty)
@@ -860,12 +885,12 @@ public class OllamaOptions implements FunctionCallingOptions, ChatOptions, Embed
@Override
public int hashCode() {
return Objects.hash(this.model, this.format, this.keepAlive, this.useNUMA, this.numCtx, this.numBatch,
this.numGPU, this.mainGPU, lowVRAM, this.f16KV, this.logitsAll, this.vocabOnly, this.useMMap,
this.useMLock, this.numThread, this.numKeep, this.seed, this.numPredict, this.topK, this.topP, tfsZ,
this.typicalP, this.repeatLastN, this.temperature, this.repeatPenalty, this.presencePenalty,
this.frequencyPenalty, this.mirostat, this.mirostatTau, this.mirostatEta, this.penalizeNewline,
this.stop, this.functionCallbacks, this.functions);
return Objects.hash(this.model, this.format, this.keepAlive, this.truncate, this.useNUMA, this.numCtx,
this.numBatch, this.numGPU, this.mainGPU, lowVRAM, this.f16KV, this.logitsAll, this.vocabOnly,
this.useMMap, this.useMLock, this.numThread, this.numKeep, this.seed, this.numPredict, this.topK,
this.topP, tfsZ, this.typicalP, this.repeatLastN, this.temperature, this.repeatPenalty,
this.presencePenalty, this.frequencyPenalty, this.mirostat, this.mirostatTau, this.mirostatEta,
this.penalizeNewline, this.stop, this.functionCallbacks, this.functions);
}
}

View File

@@ -27,7 +27,7 @@ import org.springframework.ai.ollama.api.OllamaModel;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.ollama.OllamaContainer;
import org.springframework.ai.embedding.EmbeddingRequest;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaApiIT;
@@ -66,11 +66,17 @@ class OllamaEmbeddingModelIT {
private OllamaEmbeddingModel embeddingModel;
@Test
void singleEmbedding() {
void embeddings() {
assertThat(embeddingModel).isNotNull();
EmbeddingResponse embeddingResponse = embeddingModel.embedForResponse(List.of("Hello World"));
assertThat(embeddingResponse.getResults()).hasSize(1);
EmbeddingResponse embeddingResponse = embeddingModel.call(new EmbeddingRequest(
List.of("Hello World", "Something else"), OllamaOptions.builder().withTruncate(false).build()));
assertThat(embeddingResponse.getResults()).hasSize(2);
assertThat(embeddingResponse.getResults().get(0).getIndex()).isEqualTo(0);
assertThat(embeddingResponse.getResults().get(0).getOutput()).isNotEmpty();
assertThat(embeddingResponse.getResults().get(1).getIndex()).isEqualTo(1);
assertThat(embeddingResponse.getResults().get(1).getOutput()).isNotEmpty();
assertThat(embeddingResponse.getMetadata().getModel()).isEqualTo(MODEL);
assertThat(embeddingModel.dimensions()).isEqualTo(4096);
}

View File

@@ -0,0 +1,112 @@
/*
* Copyright 2024 - 2024 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.ollama;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.Mockito.when;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.ArgumentCaptor;
import org.mockito.Captor;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.springframework.ai.embedding.EmbeddingOptions;
import org.springframework.ai.embedding.EmbeddingRequest;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.embedding.EmbeddingResultMetadata;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaApi.EmbeddingsRequest;
import org.springframework.ai.ollama.api.OllamaApi.EmbeddingsResponse;
import org.springframework.ai.ollama.api.OllamaOptions;
/**
* @author Christian Tzolov
* @since 1.0.0
*/
@ExtendWith(MockitoExtension.class)
public class OllamaEmbeddingModelTests {
@Mock
OllamaApi ollamaApi;
@Captor
ArgumentCaptor<EmbeddingsRequest> embeddingsRequestCaptor;
@Test
public void options() {
when(ollamaApi.embed(embeddingsRequestCaptor.capture()))
.thenReturn(
new EmbeddingsResponse("RESPONSE_MODEL_NAME", List.of(List.of(1d, 2d, 3d), List.of(4d, 5d, 6d))))
.thenReturn(new EmbeddingsResponse("RESPONSE_MODEL_NAME2",
List.of(List.of(7d, 8d, 9d), List.of(10d, 11d, 12d))));
// Tests default options
var defaultOptions = OllamaOptions.builder().withModel("DEFAULT_MODEL").build();
var embeddingModel = new OllamaEmbeddingModel(ollamaApi, defaultOptions);
EmbeddingResponse response = embeddingModel
.call(new EmbeddingRequest(List.of("Input1", "Input2", "Input3"), EmbeddingOptions.EMPTY));
assertThat(response.getResults()).hasSize(2);
assertThat(response.getResults().get(0).getIndex()).isEqualTo(0);
assertThat(response.getResults().get(0).getOutput()).isEqualTo(List.of(1d, 2d, 3d));
assertThat(response.getResults().get(0).getMetadata()).isEqualTo(EmbeddingResultMetadata.EMPTY);
assertThat(response.getResults().get(1).getIndex()).isEqualTo(1);
assertThat(response.getResults().get(1).getOutput()).isEqualTo(List.of(4d, 5d, 6d));
assertThat(response.getResults().get(1).getMetadata()).isEqualTo(EmbeddingResultMetadata.EMPTY);
assertThat(response.getMetadata().getModel()).isEqualTo("RESPONSE_MODEL_NAME");
assertThat(embeddingsRequestCaptor.getValue().keepAlive()).isNull();
assertThat(embeddingsRequestCaptor.getValue().truncate()).isNull();
assertThat(embeddingsRequestCaptor.getValue().input()).isEqualTo(List.of("Input1", "Input2", "Input3"));
assertThat(embeddingsRequestCaptor.getValue().options()).isEqualTo(Map.of());
assertThat(embeddingsRequestCaptor.getValue().model()).isEqualTo("DEFAULT_MODEL");
// Tests runtime options
var runtimeOptions = OllamaOptions.builder()
.withModel("RUNTIME_MODEL")
.withKeepAlive("10m")
.withTruncate(false)
.withMainGPU(666)
.build();
response = embeddingModel.call(new EmbeddingRequest(List.of("Input4", "Input5", "Input6"), runtimeOptions));
assertThat(response.getResults()).hasSize(2);
assertThat(response.getResults().get(0).getIndex()).isEqualTo(0);
assertThat(response.getResults().get(0).getOutput()).isEqualTo(List.of(7d, 8d, 9d));
assertThat(response.getResults().get(0).getMetadata()).isEqualTo(EmbeddingResultMetadata.EMPTY);
assertThat(response.getResults().get(1).getIndex()).isEqualTo(1);
assertThat(response.getResults().get(1).getOutput()).isEqualTo(List.of(10d, 11d, 12d));
assertThat(response.getResults().get(1).getMetadata()).isEqualTo(EmbeddingResultMetadata.EMPTY);
assertThat(response.getMetadata().getModel()).isEqualTo("RESPONSE_MODEL_NAME2");
assertThat(embeddingsRequestCaptor.getValue().keepAlive()).isEqualTo(Duration.ofMinutes(10));
assertThat(embeddingsRequestCaptor.getValue().truncate()).isFalse();
assertThat(embeddingsRequestCaptor.getValue().input()).isEqualTo(List.of("Input4", "Input5", "Input6"));
assertThat(embeddingsRequestCaptor.getValue().options()).isEqualTo(Map.of("main_gpu", 666));
assertThat(embeddingsRequestCaptor.getValue().model()).isEqualTo("RUNTIME_MODEL");
}
}

View File

@@ -21,6 +21,8 @@ import org.springframework.ai.ollama.api.OllamaOptions;
import static org.assertj.core.api.Assertions.assertThat;
import java.util.List;
/**
* @author Christian Tzolov
*/
@@ -32,13 +34,13 @@ public class OllamaEmbeddingRequestTests {
@Test
public void ollamaEmbeddingRequestDefaultOptions() {
var request = chatModel.ollamaEmbeddingRequest("Hello", null);
var request = chatModel.ollamaEmbeddingRequest(List.of("Hello"), null);
assertThat(request.model()).isEqualTo("DEFAULT_MODEL");
assertThat(request.options().get("num_gpu")).isEqualTo(1);
assertThat(request.options().get("main_gpu")).isEqualTo(11);
assertThat(request.options().get("use_mmap")).isEqualTo(true);
assertThat(request.prompt()).isEqualTo("Hello");
assertThat(request.input()).isEqualTo(List.of("Hello"));
}
@Test
@@ -50,13 +52,13 @@ public class OllamaEmbeddingRequestTests {
.withUseMMap(true)//
.withNumGPU(2);
var request = chatModel.ollamaEmbeddingRequest("Hello", promptOptions);
var request = chatModel.ollamaEmbeddingRequest(List.of("Hello"), promptOptions);
assertThat(request.model()).isEqualTo("PROMPT_MODEL");
assertThat(request.options().get("num_gpu")).isEqualTo(2);
assertThat(request.options().get("main_gpu")).isEqualTo(22);
assertThat(request.options().get("use_mmap")).isEqualTo(true);
assertThat(request.prompt()).isEqualTo("Hello");
assertThat(request.input()).isEqualTo(List.of("Hello"));
}
}

View File

@@ -15,6 +15,8 @@
*/
package org.springframework.ai.ollama.api;
import static org.assertj.core.api.Assertions.assertThat;
import java.io.IOException;
import java.util.List;
import java.util.stream.Collectors;
@@ -25,21 +27,19 @@ import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Test;
import org.springframework.ai.ollama.OllamaImage;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.ollama.OllamaContainer;
import reactor.core.publisher.Flux;
import org.springframework.ai.ollama.api.OllamaApi.ChatRequest;
import org.springframework.ai.ollama.api.OllamaApi.ChatResponse;
import org.springframework.ai.ollama.api.OllamaApi.EmbeddingRequest;
import org.springframework.ai.ollama.api.OllamaApi.EmbeddingResponse;
import org.springframework.ai.ollama.api.OllamaApi.EmbeddingsRequest;
import org.springframework.ai.ollama.api.OllamaApi.EmbeddingsResponse;
import org.springframework.ai.ollama.api.OllamaApi.GenerateRequest;
import org.springframework.ai.ollama.api.OllamaApi.GenerateResponse;
import org.springframework.ai.ollama.api.OllamaApi.Message;
import org.springframework.ai.ollama.api.OllamaApi.Message.Role;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.ollama.OllamaContainer;
import static org.assertj.core.api.Assertions.assertThat;;
import reactor.core.publisher.Flux;;
/**
* @author Christian Tzolov
@@ -145,12 +145,13 @@ public class OllamaApiIT {
@Test
public void embedText() {
EmbeddingRequest request = new EmbeddingRequest(MODEL, "I like to eat apples");
EmbeddingsRequest request = new EmbeddingsRequest(MODEL, "I like to eat apples");
EmbeddingResponse response = ollamaApi.embeddings(request);
EmbeddingsResponse response = ollamaApi.embed(request);
assertThat(response).isNotNull();
assertThat(response.embedding()).hasSize(3200);
assertThat(response.embeddings()).hasSize(1);
assertThat(response.embeddings().get(0)).hasSize(3200);
}
}

View File

@@ -57,21 +57,21 @@ public class Embedding implements ModelResult<List<Double>> {
*/
@Override
public List<Double> getOutput() {
return embedding;
return this.embedding;
}
/**
* @return Get the embedding index in a list of embeddings.
*/
public Integer getIndex() {
return index;
return this.index;
}
/**
* @return Get the metadata associated with the embedding.
*/
public EmbeddingResultMetadata getMetadata() {
return metadata;
return this.metadata;
}
@Override
@@ -80,8 +80,8 @@ public class Embedding implements ModelResult<List<Double>> {
return true;
if (o == null || getClass() != o.getClass())
return false;
Embedding embedding1 = (Embedding) o;
return Objects.equals(embedding, embedding1.embedding) && Objects.equals(index, embedding1.index);
Embedding other = (Embedding) o;
return Objects.equals(this.embedding, other.embedding) && Objects.equals(this.index, other.index);
}
@Override
@@ -92,7 +92,7 @@ public class Embedding implements ModelResult<List<Double>> {
@Override
public String toString() {
String message = this.embedding.isEmpty() ? "<empty>" : "<has data>";
return "Embedding{" + "embedding=" + message + ", index=" + index + '}';
return "Embedding{" + "embedding=" + message + ", index=" + this.index + '}';
}
}

View File

@@ -15,19 +15,17 @@
*/
package org.springframework.ai.transformer.splitter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.ContentFormatter;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentTransformer;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.ContentFormatter;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentTransformer;
public abstract class TextSplitter implements DocumentTransformer {
private static final Logger logger = LoggerFactory.getLogger(TextSplitter.class);

View File

@@ -1,27 +1,17 @@
= Ollama Embeddings
With https://ollama.ai/[Ollama] you can run various Large Language Models (LLMs) locally and generate embeddings from them.
Spring AI supports the Ollama text embeddings with `OllamaEmbeddingModel`.
With https://ollama.ai/[Ollama] you can run various https://ollama.com/search?c=embedding[AI Models] locally and generate embeddings from them.
An embedding is a vector (list) of floating point numbers.
The distance between two vectors measures their relatedness.
Small distances suggest high relatedness and large distances suggest low relatedness.
The `OllamaEmbeddingModel` implementation lerverages the Ollama https://github.com/ollama/ollama/blob/main/docs/api.md#generate-embeddings[Embeddings API] endpoint.
== Prerequisites
You first need to run Ollama on your local machine.
Refer to the official Ollama project link:https://github.com/ollama/ollama[README] to get started running models on your local machine.
NOTE: installing `ollama run llama3` will download a 4.7GB model artifact.
=== Add Repositories and BOM
Spring AI artifacts are published in Spring Milestone and Snapshot repositories. Refer to the xref:getting-started.adoc#repositories[Repositories] section to add these repositories to your build system.
To help with dependency management, Spring AI provides a BOM (bill of materials) to ensure that a consistent version of Spring AI is used throughout the entire project. Refer to the xref:getting-started.adoc#dependency-management[Dependency Management] section to add the Spring AI BOM to your build system.
== Auto-configuration
Spring AI provides Spring Boot auto-configuration for the Azure Ollama Embedding Model.
@@ -45,65 +35,70 @@ dependencies {
----
TIP: Refer to the xref:getting-started.adoc#dependency-management[Dependency Management] section to add the Spring AI BOM to your build file.
The `spring.ai.ollama.embedding.options.*` properties are used to configure the default options used for all embedding requests.
(It is used as `OllamaEmbeddingModel#withDefaultOptions()` instance).
Spring AI artifacts are published in Spring Milestone and Snapshot repositories. Refer to the Repositories section to add these repositories to your build system.
=== Embedding Properties
The prefix `spring.ai.ollama` is the property prefix to configure the connection to Ollama
[cols="3,6,2"]
[cols="3,6,1"]
|====
| Property | Description | Default
| spring.ai.ollama.base-url | Base URL where Ollama API server is running. | `http://localhost:11434`
|====
The prefix `spring.ai.ollama.embedding.options` is the property prefix that configures the `EmbeddingModel` implementation for Ollama.
The prefix `spring.ai.ollama.embedding.options` is the property prefix that configures the Ollama embedding model .
It includes the Ollama request (advanced) parameters such as the `model`, `keep-alive`, and `truncate` as well as the Ollama model `options` properties.
[cols="3,5,1"]
Here are the advanced request parameter for the Ollama embedding model:
[cols="3,6,1"]
|====
| Property | Description | Default
| spring.ai.ollama.embedding.enabled | Enable Ollama embedding model. | true
| spring.ai.ollama.embedding.options.model | The name of the https://github.com/ollama/ollama?tab=readme-ov-file#model-library[supported model] to use. | mistral
| spring.ai.ollama.embedding.enabled | Enables the Ollama embedding model auto-configuration. | true
| spring.ai.ollama.embedding.options.model | The name of the https://github.com/ollama/ollama?tab=readme-ov-file#model-library[supported model] to use.
You can use dedicated https://ollama.com/search?c=embedding[Embedding Model] types | mistral
| spring.ai.ollama.embedding.options.keep_alive | Controls how long the model will stay loaded into memory following the request | 5m
| spring.ai.ollama.embedding.options.truncate | Truncates the end of each input to fit within context length. Returns error if false and context length is exceeded. | true
|====
The remaining `options` properties are based on the link:https://github.com/ollama/ollama/blob/main/docs/modelfile.md#valid-parameters-and-values[Ollama Valid Parameters and Values] and link:https://github.com/ollama/ollama/blob/main/api/types.go[Ollama Types]. The default values are based on: link:https://github.com/ollama/ollama/blob/b538dc3858014f94b099730a592751a5454cab0a/api/types.go#L364[Ollama type defaults].
[cols="3,5,1"]
[cols="3,6,1"]
|====
| Property | Description | Default
| spring.ai.ollama.embedding.options.numa | Whether to use NUMA. | false
| spring.ai.ollama.embedding.options.num-ctx | Sets the size of the context window used to generate the next token. | 2048
| spring.ai.ollama.embedding.options.num-batch | ??? | -
| spring.ai.ollama.embedding.options.num-gpu | The number of layers to send to the GPU(s). On macOS it defaults to 1 to enable metal support, 0 to disable. | -
| spring.ai.ollama.embedding.options.main-gpu | ??? | -
| spring.ai.ollama.embedding.options.low-vram | ??? | -
| spring.ai.ollama.embedding.options.f16-kv | ??? | -
| spring.ai.ollama.embedding.options.logits-all | ??? | -
| spring.ai.ollama.embedding.options.vocab-only | ??? | -
| spring.ai.ollama.embedding.options.use-mmap | ??? | -
| spring.ai.ollama.embedding.options.use-mlock | ??? | -
| spring.ai.ollama.embedding.options.num-thread | Sets the number of threads to use during computation. By default, Ollama will detect this for optimal performance. It is recommended to set this value to the number of physical CPU cores your system has (as opposed to the logical number of cores). | -
| spring.ai.ollama.embedding.options.num-keep | ??? | -
| spring.ai.ollama.embedding.options.seed | Sets the random number seed to use for generation. Setting this to a specific number will make the model generate the same text for the same prompt. | 0
| spring.ai.ollama.embedding.options.num-predict | Maximum number of tokens to predict when generating text. (Default: 128, -1 = infinite generation, -2 = fill context) | 128
| spring.ai.ollama.embedding.options.num-batch | Prompt processing maximum batch size. | 512
| spring.ai.ollama.embedding.options.num-gpu | The number of layers to send to the GPU(s). On macOS it defaults to 1 to enable metal support, 0 to disable. 1 here indicates that NumGPU should be set dynamically | -1
| spring.ai.ollama.embedding.options.main-gpu | When using multiple GPUs this option controls which GPU is used for small tensors for which the overhead of splitting the computation across all GPUs is not worthwhile. The GPU in question will use slightly more VRAM to store a scratch buffer for temporary results. | 0
| spring.ai.ollama.embedding.options.low-vram | - | false
| spring.ai.ollama.embedding.options.f16-kv | - | true
| spring.ai.ollama.embedding.options.logits-all | Return logits for all the tokens, not just the last one. To enable completions to return logprobs, this must be true. | -
| spring.ai.ollama.embedding.options.vocab-only | Load only the vocabulary, not the weights. | -
| spring.ai.ollama.embedding.options.use-mmap | By default, models are mapped into memory, which allows the system to load only the necessary parts of the model as needed. However, if the model is larger than your total amount of RAM or if your system is low on available memory, using mmap might increase the risk of pageouts, negatively impacting performance. Disabling mmap results in slower load times but may reduce pageouts if you're not using mlock. Note that if the model is larger than the total amount of RAM, turning off mmap would prevent the model from loading at all. | null
| spring.ai.ollama.embedding.options.use-mlock | Lock the model in memory, preventing it from being swapped out when memory-mapped. This can improve performance but trades away some of the advantages of memory-mapping by requiring more RAM to run and potentially slowing down load times as the model loads into RAM. | false
| spring.ai.ollama.embedding.options.num-thread | Sets the number of threads to use during computation. By default, Ollama will detect this for optimal performance. It is recommended to set this value to the number of physical CPU cores your system has (as opposed to the logical number of cores). 0 = let the runtime decide | 0
| spring.ai.ollama.embedding.options.num-keep | - | 4
| spring.ai.ollama.embedding.options.seed | Sets the random number seed to use for generation. Setting this to a specific number will make the model generate the same text for the same prompt. | -1
| spring.ai.ollama.embedding.options.num-predict | Maximum number of tokens to predict when generating text. (-1 = infinite generation, -2 = fill context) | -1
| spring.ai.ollama.embedding.options.top-k | Reduces the probability of generating nonsense. A higher value (e.g., 100) will give more diverse answers, while a lower value (e.g., 10) will be more conservative. | 40
| spring.ai.ollama.embedding.options.top-p | Works together with top-k. A higher value (e.g., 0.95) will lead to more diverse text, while a lower value (e.g., 0.5) will generate more focused and conservative text. | 0.9
| spring.ai.ollama.embedding.options.tfs-z | Tail-free sampling is used to reduce the impact of less probable tokens from the output. A higher value (e.g., 2.0) will reduce the impact more, while a value of 1.0 disables this setting. | 1
| spring.ai.ollama.embedding.options.typical-p | ??? | -
| spring.ai.ollama.embedding.options.repeat-last-n | Sets how far back for the model to look back to prevent repetition. (Default: 64, 0 = disabled, -1 = num_ctx) | 64
| spring.ai.ollama.embedding.options.tfs-z | Tail-free sampling is used to reduce the impact of less probable tokens from the output. A higher value (e.g., 2.0) will reduce the impact more, while a value of 1.0 disables this setting. | 1.0
| spring.ai.ollama.embedding.options.typical-p | - | 1.0
| spring.ai.ollama.embedding.options.repeat-last-n | Sets how far back for the model to look back to prevent repetition. (Default: 64, 0 = disabled, -1 = num_ctx) | 64
| spring.ai.ollama.embedding.options.temperature | The temperature of the model. Increasing the temperature will make the model answer more creatively. | 0.8
| spring.ai.ollama.embedding.options.repeat-penalty | Sets how strongly to penalize repetitions. A higher value (e.g., 1.5) will penalize repetitions more strongly, while a lower value (e.g., 0.9) will be more lenient. | 1.1
| spring.ai.ollama.embedding.options.presence-penalty | ??? | -
| spring.ai.ollama.embedding.options.frequency-penalty | ??? | -
| spring.ai.ollama.embedding.options.presence-penalty | - | 0.0
| spring.ai.ollama.embedding.options.frequency-penalty | - | 0.0
| spring.ai.ollama.embedding.options.mirostat | Enable Mirostat sampling for controlling perplexity. (default: 0, 0 = disabled, 1 = Mirostat, 2 = Mirostat 2.0) | 0
| spring.ai.ollama.embedding.options.mirostat-tau | Controls the balance between coherence and diversity of the output. A lower value will result in more focused and coherent text. | 5.0
| spring.ai.ollama.embedding.options.mirostat-eta | Influences how quickly the algorithm responds to feedback from the generated text. A lower learning rate will result in slower adjustments, while a higher learning rate will make the algorithm more responsive. | 0.1
| spring.ai.ollama.embedding.options.penalize-newline | ??? | -
| spring.ai.ollama.embedding.options.penalize-newline | - | true
| spring.ai.ollama.embedding.options.stop | Sets the stop sequences to use. When this pattern is encountered the LLM will stop generating text and return. Multiple stop patterns may be set by specifying multiple separate stop parameters in a modelfile. | -
| spring.ai.ollama.embedding.options.functions | List of functions, identified by their names, to enable for function calling in a single prompt requests. Functions with those names must exist in the functionCallbacks registry. | -
|====
TIP: All properties prefixed with `spring.ai.ollama.embedding.options` can be overridden at runtime by adding a request specific <<embedding-options>> to the `EmbeddingRequest` call.
@@ -114,7 +109,7 @@ The https://github.com/spring-projects/spring-ai/blob/main/models/spring-ai-olla
The default options can be configured using the `spring.ai.ollama.embedding.options` properties as well.
At start-time use the `OllamaEmbeddingModel#withDefaultOptions()` to configure the default options used for all embedding requests.
At start-time use the `OllamaEmbeddingModel(OllamaApi ollamaApi, OllamaOptions defaultOptions)` to configure the default options used for all embedding requests.
At run-time you can override the default options, using a `OllamaOptions` instance as part of your `EmbeddingRequest`.
For example to override the default model name for a specific request:
@@ -123,9 +118,10 @@ For example to override the default model name for a specific request:
----
EmbeddingResponse embeddingResponse = embeddingModel.call(
new EmbeddingRequest(List.of("Hello World", "World is big and salvation is near"),
OllamaOptions.create()
OllamaOptions.builder()
.withModel("Different-Embedding-Model-Deployment-Name"))
);
.withtTuncates(false)
.build());
----
== Sample Controller
@@ -180,19 +176,24 @@ TIP: Refer to the xref:getting-started.adoc#dependency-management[Dependency Man
NOTE: The `spring-ai-ollama` dependency provides access also to the `OllamaChatModel`.
For more information about the `OllamaChatModel` refer to the link:../chat/ollama-chat.html[Ollama Chat Client] section.
Next, create an `OllamaEmbeddingModel` instance and use it to compute the similarity between two input texts:
Next, create an `OllamaEmbeddingModel` instance and use it to compute the embeddings for two input texts using a dedicated `chroma/all-minilm-l6-v2-f32` embedding models:
[source,java]
----
var ollamaApi = new OllamaApi();
var embeddingModel = new OllamaEmbeddingModel(ollamaApi)
.withDefaultOptions(OllamaOptions.create()
.withModel(OllamaOptions.DEFAULT_MODEL)
var embeddingModel = new OllamaEmbeddingModel(ollamaApi,
OllamaOptions.builder()
.withModel(OllamaModel.MISTRAL.id())
.build()
.toMap());
EmbeddingResponse embeddingResponse = embeddingModel
.embedForResponse(List.of("Hello World", "World is big and salvation is near"));
EmbeddingResponse embeddingResponse = embeddingModel.call(
new EmbeddingRequest(List.of("Hello World", "World is big and salvation is near"),
OllamaOptions.builder()
.withModel("chroma/all-minilm-l6-v2-f32"))
.withtTruncate(false)
.build());
----
The `OllamaOptions` provides the configuration information for all embedding requests.