聊聊Spring AI的EmbeddingModel

本文主要研究一下Spring AI的EmbeddingModel

EmbeddingModel

spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingModel.java

public interface EmbeddingModel extends Model<EmbeddingRequest, EmbeddingResponse> {

    @Override
    EmbeddingResponse call(EmbeddingRequest request);

    /**
     * Embeds the given text into a vector.
     * @param text the text to embed.
     * @return the embedded vector.
     */
    default float[] embed(String text) {
        Assert.notNull(text, "Text must not be null");
        List<float[]> response = this.embed(List.of(text));
        return response.iterator().next();
    }

    /**
     * Embeds the given document's content into a vector.
     * @param document the document to embed.
     * @return the embedded vector.
     */
    float[] embed(Document document);

    /**
     * Embeds a batch of texts into vectors.
     * @param texts list of texts to embed.
     * @return list of embedded vectors.
     */
    default List<float[]> embed(List<String> texts) {
        Assert.notNull(texts, "Texts must not be null");
        return this.call(new EmbeddingRequest(texts, EmbeddingOptionsBuilder.builder().build()))
            .getResults()
            .stream()
            .map(Embedding::getOutput)
            .toList();
    }

    /**
     * Embeds a batch of {@link Document}s into vectors based on a
     * {@link BatchingStrategy}.
     * @param documents list of {@link Document}s.
     * @param options {@link EmbeddingOptions}.
     * @param batchingStrategy {@link BatchingStrategy}.
     * @return a list of float[] that represents the vectors for the incoming
     * {@link Document}s. The returned list is expected to be in the same order of the
     * {@link Document} list.
     */
    default List<float[]> embed(List<Document> documents, EmbeddingOptions options, BatchingStrategy batchingStrategy) {
        Assert.notNull(documents, "Documents must not be null");
        List<float[]> embeddings = new ArrayList<>(documents.size());
        List<List<Document>> batch = batchingStrategy.batch(documents);
        for (List<Document> subBatch : batch) {
            List<String> texts = subBatch.stream().map(Document::getText).toList();
            EmbeddingRequest request = new EmbeddingRequest(texts, options);
            EmbeddingResponse response = this.call(request);
            for (int i = 0; i < subBatch.size(); i++) {
                embeddings.add(response.getResults().get(i).getOutput());
            }
        }
        Assert.isTrue(embeddings.size() == documents.size(),
                "Embeddings must have the same number as that of the documents");
        return embeddings;
    }

    /**
     * Embeds a batch of texts into vectors and returns the {@link EmbeddingResponse}.
     * @param texts list of texts to embed.
     * @return the embedding response.
     */
    default EmbeddingResponse embedForResponse(List<String> texts) {
        Assert.notNull(texts, "Texts must not be null");
        return this.call(new EmbeddingRequest(texts, EmbeddingOptionsBuilder.builder().build()));
    }

    /**
     * Get the number of dimensions of the embedded vectors. Note that by default, this
     * method will call the remote Embedding endpoint to get the dimensions of the
     * embedded vectors. If the dimensions are known ahead of time, it is recommended to
     * override this method.
     * @return the number of dimensions of the embedded vectors.
     */
    default int dimensions() {
        return embed("Test String").length;
    }

}

EmbeddingModel继承了Model接口,其入参类型为EmbeddingRequest,返回类型为EmbeddingResponse,它定义了call、embed接口,提供了embed、embedForResponse、dimensions的默认实现

EmbeddingRequest

spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingRequest.java

public class EmbeddingRequest implements ModelRequest<List<String>> {

    private final List<String> inputs;

    private final EmbeddingOptions options;

    public EmbeddingRequest(List<String> inputs, EmbeddingOptions options) {
        this.inputs = inputs;
        this.options = options;
    }

    @Override
    public List<String> getInstructions() {
        return this.inputs;
    }

    @Override
    public EmbeddingOptions getOptions() {
        return this.options;
    }

}

EmbeddingRequest实现了ModelRequest接口,其getInstructions返回的是List<String>

EmbeddingResponse

spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingResponse.java

public class EmbeddingResponse implements ModelResponse<Embedding> {

    /**
     * Embedding data.
     */
    private final List<Embedding> embeddings;

    /**
     * Embedding metadata.
     */
    private final EmbeddingResponseMetadata metadata;

    /**
     * Creates a new {@link EmbeddingResponse} instance with empty metadata.
     * @param embeddings the embedding data.
     */
    public EmbeddingResponse(List<Embedding> embeddings) {
        this(embeddings, new EmbeddingResponseMetadata());
    }

    /**
     * Creates a new {@link EmbeddingResponse} instance.
     * @param embeddings the embedding data.
     * @param metadata the embedding metadata.
     */
    public EmbeddingResponse(List<Embedding> embeddings, EmbeddingResponseMetadata metadata) {
        this.embeddings = embeddings;
        this.metadata = metadata;
    }

    /**
     * @return Get the embedding metadata.
     */
    public EmbeddingResponseMetadata getMetadata() {
        return this.metadata;
    }

    @Override
    public Embedding getResult() {
        Assert.notEmpty(this.embeddings, "No embedding data available.");
        return this.embeddings.get(0);
    }

    /**
     * @return Get the embedding data.
     */
    @Override
    public List<Embedding> getResults() {
        return this.embeddings;
    }

    //......
}   

EmbeddingResponse实现了ModelResponse接口,其result为Embedding类型

AbstractEmbeddingModel

spring-ai-core/src/main/java/org/springframework/ai/embedding/AbstractEmbeddingModel.java

public abstract class AbstractEmbeddingModel implements EmbeddingModel {

    private static final Map<String, Integer> KNOWN_EMBEDDING_DIMENSIONS = loadKnownModelDimensions();

    /**
     * Default constructor.
     */
    public AbstractEmbeddingModel() {
    }

    /**
     * Cached embedding dimensions.
     */
    protected final AtomicInteger embeddingDimensions = new AtomicInteger(-1);

    /**
     * Return the dimension of the requested embedding generative name. If the generative
     * name is unknown uses the EmbeddingModel to perform a dummy EmbeddingModel#embed and
     * count the response dimensions.
     * @param embeddingModel Fall-back client to determine, empirically the dimensions.
     * @param modelName Embedding generative name to retrieve the dimensions for.
     * @param dummyContent Dummy content to use for the empirical dimension calculation.
     * @return Returns the embedding dimensions for the modelName.
     */
    public static int dimensions(EmbeddingModel embeddingModel, String modelName, String dummyContent) {

        if (KNOWN_EMBEDDING_DIMENSIONS.containsKey(modelName)) {
            // Retrieve the dimension from a pre-configured file.
            return KNOWN_EMBEDDING_DIMENSIONS.get(modelName);
        }
        else {
            // Determine the dimensions empirically.
            // Generate an embedding and count the dimension size;
            return embeddingModel.embed(dummyContent).length;
        }
    }

    private static Map<String, Integer> loadKnownModelDimensions() {
        try {
            Properties properties = new Properties();
            properties.load(new DefaultResourceLoader()
                .getResource("classpath:/embedding/embedding-model-dimensions.properties")
                .getInputStream());
            return properties.entrySet()
                .stream()
                .collect(Collectors.toMap(e -> e.getKey().toString(), e -> Integer.parseInt(e.getValue().toString())));
        }
        catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public int dimensions() {
        if (this.embeddingDimensions.get() < 0) {
            this.embeddingDimensions.set(dimensions(this, "Test", "Hello World"));
        }
        return this.embeddingDimensions.get();
    }

}

AbstractEmbeddingModel实现了EmbeddingModel接口定义的dimensions方法,它在不同模块有不同的实现子类,比如spring-ai-openai的OpenAiEmbeddingModel、spring-ai-ollama的OllamaEmbeddingModel、spring-ai-minimax的MiniMaxEmbeddingModel

OllamaEmbeddingAutoConfiguration

org/springframework/ai/model/ollama/autoconfigure/OllamaEmbeddingAutoConfiguration.java

@AutoConfiguration(after = RestClientAutoConfiguration.class)
@ConditionalOnClass(OllamaEmbeddingModel.class)
@ConditionalOnProperty(name = SpringAIModelProperties.EMBEDDING_MODEL, havingValue = SpringAIModels.OLLAMA,
        matchIfMissing = true)
@EnableConfigurationProperties({ OllamaEmbeddingProperties.class, OllamaInitializationProperties.class })
@ImportAutoConfiguration(classes = { OllamaApiAutoConfiguration.class, RestClientAutoConfiguration.class,
        WebClientAutoConfiguration.class })
public class OllamaEmbeddingAutoConfiguration {

    @Bean
    @ConditionalOnMissingBean
    public OllamaEmbeddingModel ollamaEmbeddingModel(OllamaApi ollamaApi, OllamaEmbeddingProperties properties,
            OllamaInitializationProperties initProperties, ObjectProvider<ObservationRegistry> observationRegistry,
            ObjectProvider<EmbeddingModelObservationConvention> observationConvention) {
        var embeddingModelPullStrategy = initProperties.getEmbedding().isInclude()
                ? initProperties.getPullModelStrategy() : PullModelStrategy.NEVER;

        var embeddingModel = OllamaEmbeddingModel.builder()
            .ollamaApi(ollamaApi)
            .defaultOptions(properties.getOptions())
            .observationRegistry(observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP))
            .modelManagementOptions(new ModelManagementOptions(embeddingModelPullStrategy,
                    initProperties.getEmbedding().getAdditionalModels(), initProperties.getTimeout(),
                    initProperties.getMaxRetries()))
            .build();

        observationConvention.ifAvailable(embeddingModel::setObservationConvention);

        return embeddingModel;
    }

}

OllamaEmbeddingAutoConfiguration在spring.ai.model.embeddingollama时启用,它自动配置了OllamaEmbeddingModel

OllamaEmbeddingProperties

org/springframework/ai/model/ollama/autoconfigure/OllamaEmbeddingProperties.java

@ConfigurationProperties(OllamaEmbeddingProperties.CONFIG_PREFIX)
public class OllamaEmbeddingProperties {

    public static final String CONFIG_PREFIX = "spring.ai.ollama.embedding";

    /**
     * Client lever Ollama options. Use this property to configure generative temperature,
     * topK and topP and alike parameters. The null values are ignored defaulting to the
     * generative's defaults.
     */
    @NestedConfigurationProperty
    private OllamaOptions options = OllamaOptions.builder().model(OllamaModel.MXBAI_EMBED_LARGE.id()).build();

    public String getModel() {
        return this.options.getModel();
    }

    public void setModel(String model) {
        this.options.setModel(model);
    }

    public OllamaOptions getOptions() {
        return this.options;
    }

}

OllamaEmbeddingProperties主要是提供了OllamaOptions属性配置,具体可以参考https://github.com/ggerganov/llama.cpp/blob/master/examples/main/README.md

OllamaInitializationProperties

org/springframework/ai/model/ollama/autoconfigure/OllamaInitializationProperties.java

@ConfigurationProperties(OllamaInitializationProperties.CONFIG_PREFIX)
public class OllamaInitializationProperties {

    public static final String CONFIG_PREFIX = "spring.ai.ollama.init";

    /**
     * Chat models initialization settings.
     */
    private final ModelTypeInit chat = new ModelTypeInit();

    /**
     * Embedding models initialization settings.
     */
    private final ModelTypeInit embedding = new ModelTypeInit();

    /**
     * Whether to pull models at startup-time and how.
     */
    private PullModelStrategy pullModelStrategy = PullModelStrategy.NEVER;

    /**
     * How long to wait for a model to be pulled.
     */
    private Duration timeout = Duration.ofMinutes(5);

    /**
     * Maximum number of retries for the model pull operation.
     */
    private int maxRetries = 0;

    public PullModelStrategy getPullModelStrategy() {
        return this.pullModelStrategy;
    }

    public void setPullModelStrategy(PullModelStrategy pullModelStrategy) {
        this.pullModelStrategy = pullModelStrategy;
    }

    public ModelTypeInit getChat() {
        return this.chat;
    }

    public ModelTypeInit getEmbedding() {
        return this.embedding;
    }

    public Duration getTimeout() {
        return this.timeout;
    }

    public void setTimeout(Duration timeout) {
        this.timeout = timeout;
    }

    public int getMaxRetries() {
        return this.maxRetries;
    }

    public void setMaxRetries(int maxRetries) {
        this.maxRetries = maxRetries;
    }

    public static class ModelTypeInit {

        /**
         * Include this type of models in the initialization task.
         */
        private boolean include = true;

        /**
         * Additional models to initialize besides the ones configured via default
         * properties.
         */
        private List<String> additionalModels = List.of();

        public boolean isInclude() {
            return this.include;
        }

        public void setInclude(boolean include) {
            this.include = include;
        }

        public List<String> getAdditionalModels() {
            return this.additionalModels;
        }

        public void setAdditionalModels(List<String> additionalModels) {
            this.additionalModels = additionalModels;
        }

    }

}

OllamaInitializationProperties提供了spring.ai.ollama.init即ollama初始化的相关配置,其中ModelTypeInit可以指定初始化哪些额外的model

示例

pom.xml

<dependency>
   <groupId>org.springframework.ai</groupId>
   <artifactId>spring-ai-starter-model-ollama</artifactId>
</dependency>

配置

spring:
  ai:
    model:
      embedding: ollama
    ollama:
      init:
        timeout: 5m
        max-retries: 0
        embedding:
          include: true
          additional-models: []
      base-url: http://localhost:11434
      embedding:
        enabled: true
        options:
          model: bge-m3:latest
          truncate: true

example

    @Test
    public void testCall() {
        EmbeddingRequest request = new EmbeddingRequest(List.of("Hello World", "World is big and salvation is near"),
                OllamaOptions.builder()
                        .model("bge-m3:latest")
                        .truncate(false)
                        .build());
        EmbeddingResponse embeddingResponse = embeddingModel.call(request);
        log.info("resp:{}", JSON.toJSONString(embeddingResponse));
    }

小结

Spring AI定义了EmbeddingModel接口,它继承了Model接口,其入参类型为EmbeddingRequest,返回类型为EmbeddingResponse,它定义了call、embed接口,提供了embed、embedForResponse、dimensions的默认实现;AbstractEmbeddingModel实现了EmbeddingModel接口定义的dimensions方法,它在不同模块有不同的实现子类,比如spring-ai-openai的OpenAiEmbeddingModel、spring-ai-ollama的OllamaEmbeddingModel、spring-ai-minimax的MiniMaxEmbeddingModel等;OllamaEmbeddingAutoConfiguration在spring.ai.model.embeddingollama时启用,它自动配置了OllamaEmbeddingModel。

doc

©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 226,608评论 6 524
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 97,559评论 3 411
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 174,222评论 0 371
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 62,103评论 1 306
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 70,943评论 6 405
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 54,480评论 1 318
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 42,589评论 3 433
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 41,743评论 0 283
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 48,253评论 1 329
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 40,233评论 3 352
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 42,366评论 1 363
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 37,932评论 5 354
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 43,624评论 3 342
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 34,047评论 0 25
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 35,246评论 1 278
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 50,953评论 3 385
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 47,351评论 2 369

推荐阅读更多精彩内容