From 7a347966daaacd9b427bb18928f0f984e0a9c001 Mon Sep 17 00:00:00 2001
From: Louis Dureuil <louis@meilisearch.com>
Date: Mon, 22 Jul 2024 12:09:52 +0200
Subject: [PATCH] Allow explicit `dimensions` for ollama

---
 milli/src/update/settings.rs |  2 --
 milli/src/vector/mod.rs      | 16 ----------------
 milli/src/vector/ollama.rs   | 17 ++++++++++++++---
 milli/src/vector/settings.rs | 27 ++++++++++++++++++++++-----
 4 files changed, 36 insertions(+), 26 deletions(-)

diff --git a/milli/src/update/settings.rs b/milli/src/update/settings.rs
index 2836f4bc9..e5e13a03c 100644
--- a/milli/src/update/settings.rs
+++ b/milli/src/update/settings.rs
@@ -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)?;
 
diff --git a/milli/src/vector/mod.rs b/milli/src/vector/mod.rs
index a1c937d24..caccb404b 100644
--- a/milli/src/vector/mod.rs
+++ b/milli/src/vector/mod.rs
@@ -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> {
diff --git a/milli/src/vector/ollama.rs b/milli/src/vector/ollama.rs
index d8b75342b..7d41ab4e9 100644
--- a/milli/src/vector/ollama.rs
+++ b/milli/src/vector/ollama.rs
@@ -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!({
diff --git a/milli/src/vector/settings.rs b/milli/src/vector/settings.rs
index ef0c8f7ff..2a6e18a77 100644
--- a/milli/src/vector/settings.rs
+++ b/milli/src/vector/settings.rs
@@ -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;