Allow explicit dimensions for ollama

This commit is contained in:
Louis Dureuil 2024-07-22 12:09:52 +02:00
parent 6c598fa06d
commit 7a347966da
No known key found for this signature in database
4 changed files with 36 additions and 26 deletions

View File

@ -1667,8 +1667,6 @@ pub fn validate_embedding_settings(
}
}
EmbedderSource::Ollama => {
// Dimensions get inferred, only model name is required
check_unset(&dimensions, EmbeddingSettings::DIMENSIONS, inferred_source, name)?;
check_set(&model, EmbeddingSettings::MODEL, inferred_source, name)?;
check_unset(&revision, EmbeddingSettings::REVISION, inferred_source, name)?;

View File

@ -202,22 +202,6 @@ impl Default for EmbedderOptions {
}
}
impl EmbedderOptions {
/// Default options for the Hugging Face embedder
pub fn huggingface() -> Self {
Self::HuggingFace(hf::EmbedderOptions::new())
}
/// Default options for the OpenAI embedder
pub fn openai(api_key: Option<String>) -> Self {
Self::OpenAi(openai::EmbedderOptions::with_default_model(api_key))
}
pub fn ollama(api_key: Option<String>, url: Option<String>) -> Self {
Self::Ollama(ollama::EmbedderOptions::with_default_model(api_key, url))
}
}
impl Embedder {
/// Spawns a new embedder built from its options.
pub fn new(options: EmbedderOptions) -> std::result::Result<Self, NewEmbedderError> {

View File

@ -17,11 +17,22 @@ pub struct EmbedderOptions {
pub url: Option<String>,
pub api_key: Option<String>,
pub distribution: Option<DistributionShift>,
pub dimensions: Option<usize>,
}
impl EmbedderOptions {
pub fn with_default_model(api_key: Option<String>, url: Option<String>) -> Self {
Self { embedding_model: "nomic-embed-text".into(), api_key, url, distribution: None }
pub fn with_default_model(
api_key: Option<String>,
url: Option<String>,
dimensions: Option<usize>,
) -> Self {
Self {
embedding_model: "nomic-embed-text".into(),
api_key,
url,
distribution: None,
dimensions,
}
}
}
@ -31,7 +42,7 @@ impl Embedder {
let rest_embedder = match RestEmbedder::new(
RestEmbedderOptions {
api_key: options.api_key,
dimensions: None,
dimensions: options.dimensions,
distribution: options.distribution,
url: options.url.unwrap_or_else(get_ollama_path),
request: serde_json::json!({

View File

@ -152,7 +152,18 @@ impl SettingsDiff {
ReindexAction::push_action(&mut reindex_action, ReindexAction::FullReindex);
}
if dimensions.apply(new_dimensions) {
ReindexAction::push_action(&mut reindex_action, ReindexAction::FullReindex);
match source {
// regenerate on dimensions change in OpenAI since truncation is supported
Setting::Set(EmbedderSource::OpenAi) | Setting::Reset => {
ReindexAction::push_action(
&mut reindex_action,
ReindexAction::FullReindex,
);
}
// for all other embedders, the parameter is a hint that should not be able to change the result
// and so won't cause a reindex by itself.
_ => {}
}
}
if url.apply(new_url) {
match source {
@ -329,9 +340,12 @@ impl EmbeddingSettings {
Self::API_KEY => {
&[EmbedderSource::OpenAi, EmbedderSource::Ollama, EmbedderSource::Rest]
}
Self::DIMENSIONS => {
&[EmbedderSource::OpenAi, EmbedderSource::UserProvided, EmbedderSource::Rest]
}
Self::DIMENSIONS => &[
EmbedderSource::OpenAi,
EmbedderSource::UserProvided,
EmbedderSource::Ollama,
EmbedderSource::Rest,
],
Self::DOCUMENT_TEMPLATE => &[
EmbedderSource::HuggingFace,
EmbedderSource::OpenAi,
@ -377,6 +391,7 @@ impl EmbeddingSettings {
Self::DOCUMENT_TEMPLATE,
Self::URL,
Self::API_KEY,
Self::DIMENSIONS,
Self::DISTRIBUTION,
],
EmbedderSource::UserProvided => &[Self::SOURCE, Self::DIMENSIONS, Self::DISTRIBUTION],
@ -486,12 +501,13 @@ impl From<EmbeddingConfig> for EmbeddingSettings {
url,
api_key,
distribution,
dimensions,
}) => Self {
source: Setting::Set(EmbedderSource::Ollama),
model: Setting::Set(embedding_model),
revision: Setting::NotSet,
api_key: api_key.map(Setting::Set).unwrap_or_default(),
dimensions: Setting::NotSet,
dimensions: dimensions.map(Setting::Set).unwrap_or_default(),
document_template: Setting::Set(prompt.template),
url: url.map(Setting::Set).unwrap_or_default(),
request: Setting::NotSet,
@ -583,6 +599,7 @@ impl From<EmbeddingSettings> for EmbeddingConfig {
super::ollama::EmbedderOptions::with_default_model(
api_key.set(),
url.set(),
dimensions.set(),
);
if let Some(model) = model.set() {
options.embedding_model = model;