Skip to content

Commit 8ed2af4

Browse files
sblashuktzolov
authored andcommitted
Add Ollama enum with supported models and their ids
- Add missing license header
1 parent bcb559a commit 8ed2af4

File tree

2 files changed

+96
-1
lines changed

2 files changed

+96
-1
lines changed
Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
/*
2+
* Copyright 2023 - 2024 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package org.springframework.ai.ollama.api;
17+
18+
/**
19+
* Helper class for common Ollama models.
20+
*
21+
* @author Siarhei Blashuk
22+
* @since 0.8.1
23+
*/
24+
public enum OllamaModel {
25+
26+
/**
27+
* Llama 2 is a collection of language models ranging from 7B to 70B parameters.
28+
*/
29+
LLAMA2("llama2"),
30+
31+
/**
32+
* The 7B parameters model
33+
*/
34+
MISTRAL("mistral"),
35+
36+
/**
37+
* The 2.7B uncensored Dolphin model
38+
*/
39+
DOLPHIN_PHI("dolphin-phi"),
40+
41+
/**
42+
* The Phi-2 2.7B language model
43+
*/
44+
PHI("phi"),
45+
46+
/**
47+
* A fine-tuned Mistral model
48+
*/
49+
NEURAL_CHAT("neural-chat"),
50+
51+
/**
52+
* Starling-7B model
53+
*/
54+
STARLING_LM("starling-lm"),
55+
56+
/**
57+
* Code Llama is based on Llama 2 model
58+
*/
59+
CODELLAMA("codellama"),
60+
61+
/**
62+
* Orca Mini is based on Llama and Llama 2 ranging from 3 billion parameters to 70
63+
* billion
64+
*/
65+
ORCA_MINI("orca-mini"),
66+
67+
/**
68+
* Llava is a Large Language and Vision Assistant model
69+
*/
70+
LLAVA("llava"),
71+
72+
/**
73+
* Gemma is a lightweight model with 2 billion and 7 billion
74+
*/
75+
GEMMA("gemma"),
76+
77+
/**
78+
* Uncensored Llama 2 model
79+
*/
80+
LLAMA2_UNCENSORED("llama2-uncensored");
81+
82+
private final String id;
83+
84+
OllamaModel(String id) {
85+
this.id = id;
86+
}
87+
88+
public String id() {
89+
return this.id;
90+
}
91+
92+
}

models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/api/OllamaOptions.java

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@
4343
@JsonInclude(Include.NON_NULL)
4444
public class OllamaOptions implements ChatOptions, EmbeddingOptions {
4545

46-
public static final String DEFAULT_MODEL = "mistral";
46+
public static final String DEFAULT_MODEL = OllamaModel.MISTRAL.id();
4747

4848
// @formatter:off
4949
/**
@@ -243,6 +243,9 @@ public class OllamaOptions implements ChatOptions, EmbeddingOptions {
243243
*/
244244
@JsonProperty("model") private String model;
245245

246+
/**
247+
* @param model The ollama model names to use. See the {@link OllamaModel} for the common models.
248+
*/
246249
public OllamaOptions withModel(String model) {
247250
this.model = model;
248251
return this;

0 commit comments

Comments
 (0)