Add Ollama enum with supported models and their ids

- Add missing license header
This commit is contained in:
sblashuk
2024-03-01 17:33:00 +01:00
committed by Christian Tzolov
parent bcb559a82c
commit 8ed2af4115
2 changed files with 96 additions and 1 deletions

View File

@@ -0,0 +1,92 @@
/*
* Copyright 2023 - 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.api;
/**
* Helper class for common Ollama models.
*
* @author Siarhei Blashuk
* @since 0.8.1
*/
public enum OllamaModel {
/**
* Llama 2 is a collection of language models ranging from 7B to 70B parameters.
*/
LLAMA2("llama2"),
/**
* The 7B parameters model
*/
MISTRAL("mistral"),
/**
* The 2.7B uncensored Dolphin model
*/
DOLPHIN_PHI("dolphin-phi"),
/**
* The Phi-2 2.7B language model
*/
PHI("phi"),
/**
* A fine-tuned Mistral model
*/
NEURAL_CHAT("neural-chat"),
/**
* Starling-7B model
*/
STARLING_LM("starling-lm"),
/**
* Code Llama is based on Llama 2 model
*/
CODELLAMA("codellama"),
/**
* Orca Mini is based on Llama and Llama 2 ranging from 3 billion parameters to 70
* billion
*/
ORCA_MINI("orca-mini"),
/**
* Llava is a Large Language and Vision Assistant model
*/
LLAVA("llava"),
/**
* Gemma is a lightweight model with 2 billion and 7 billion
*/
GEMMA("gemma"),
/**
* Uncensored Llama 2 model
*/
LLAMA2_UNCENSORED("llama2-uncensored");
private final String id;
OllamaModel(String id) {
this.id = id;
}
public String id() {
return this.id;
}
}

View File

@@ -43,7 +43,7 @@ import org.springframework.ai.embedding.EmbeddingOptions;
@JsonInclude(Include.NON_NULL)
public class OllamaOptions implements ChatOptions, EmbeddingOptions {
public static final String DEFAULT_MODEL = "mistral";
public static final String DEFAULT_MODEL = OllamaModel.MISTRAL.id();
// @formatter:off
/**
@@ -243,6 +243,9 @@ public class OllamaOptions implements ChatOptions, EmbeddingOptions {
*/
@JsonProperty("model") private String model;
/**
* @param model The ollama model names to use. See the {@link OllamaModel} for the common models.
*/
public OllamaOptions withModel(String model) {
this.model = model;
return this;