diff --git a/meilisearch-types/src/error.rs b/meilisearch-types/src/error.rs index 2182b1836..1b54e77c0 100644 --- a/meilisearch-types/src/error.rs +++ b/meilisearch-types/src/error.rs @@ -347,6 +347,9 @@ impl ErrorCode for milli::Error { UserError::InvalidFieldForSource { .. } | UserError::MissingFieldForSource { .. } | UserError::InvalidOpenAiModel { .. } + | UserError::InvalidOpenAiModelDimensions { .. } + | UserError::InvalidOpenAiModelDimensionsMax { .. } + | UserError::InvalidSettingsDimensions { .. } | UserError::InvalidPrompt(_) => Code::InvalidSettingsEmbedders, UserError::TooManyEmbedders(_) => Code::InvalidSettingsEmbedders, UserError::InvalidPromptForEmbeddings(..) => Code::InvalidSettingsEmbedders, diff --git a/milli/src/error.rs b/milli/src/error.rs index 5a4fbc7f5..1147085dd 100644 --- a/milli/src/error.rs +++ b/milli/src/error.rs @@ -227,6 +227,22 @@ only composed of alphanumeric characters (a-z A-Z 0-9), hyphens (-) and undersco source_: crate::vector::settings::EmbedderSource, embedder_name: String, }, + #[error("`.embedders.{embedder_name}.dimensions`: Model `{model}` does not support overriding its native dimensions of {expected_dimensions}. Found {dimensions}")] + InvalidOpenAiModelDimensions { + embedder_name: String, + model: &'static str, + dimensions: usize, + expected_dimensions: usize, + }, + #[error("`.embedders.{embedder_name}.dimensions`: Model `{model}` does not support overriding its dimensions to a value higher than {max_dimensions}. Found {dimensions}")] + InvalidOpenAiModelDimensionsMax { + embedder_name: String, + model: &'static str, + dimensions: usize, + max_dimensions: usize, + }, + #[error("`.embedders.{embedder_name}.dimensions`: `dimensions` cannot be zero")] + InvalidSettingsDimensions { embedder_name: String }, } impl From for Error { diff --git a/milli/src/update/settings.rs b/milli/src/update/settings.rs index d770bcd74..a3ba42119 100644 --- a/milli/src/update/settings.rs +++ b/milli/src/update/settings.rs @@ -974,6 +974,9 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> { crate::vector::settings::EmbeddingSettings::apply_default_source( &mut setting, ); + crate::vector::settings::EmbeddingSettings::apply_default_openai_model( + &mut setting, + ); let setting = validate_embedding_settings(setting, &name)?; changed = true; new_configs.insert(name, setting); @@ -1119,6 +1122,14 @@ pub fn validate_embedding_settings( let Setting::Set(settings) = settings else { return Ok(settings) }; let EmbeddingSettings { source, model, revision, api_key, dimensions, document_template } = settings; + + if let Some(0) = dimensions.set() { + return Err(crate::error::UserError::InvalidSettingsDimensions { + embedder_name: name.to_owned(), + } + .into()); + } + let Some(inferred_source) = source.set() else { return Ok(Setting::Set(EmbeddingSettings { source, @@ -1132,14 +1143,34 @@ pub fn validate_embedding_settings( match inferred_source { EmbedderSource::OpenAi => { check_unset(&revision, "revision", inferred_source, name)?; - check_unset(&dimensions, "dimensions", inferred_source, name)?; if let Setting::Set(model) = &model { - crate::vector::openai::EmbeddingModel::from_name(model.as_str()).ok_or( - crate::error::UserError::InvalidOpenAiModel { + let model = crate::vector::openai::EmbeddingModel::from_name(model.as_str()) + .ok_or(crate::error::UserError::InvalidOpenAiModel { embedder_name: name.to_owned(), model: model.clone(), - }, - )?; + })?; + if let Setting::Set(dimensions) = dimensions { + if !model.supports_overriding_dimensions() + && dimensions != model.default_dimensions() + { + return Err(crate::error::UserError::InvalidOpenAiModelDimensions { + embedder_name: name.to_owned(), + model: model.name(), + dimensions, + expected_dimensions: model.default_dimensions(), + } + .into()); + } + if dimensions > model.default_dimensions() { + return Err(crate::error::UserError::InvalidOpenAiModelDimensionsMax { + embedder_name: name.to_owned(), + model: model.name(), + dimensions, + max_dimensions: model.default_dimensions(), + } + .into()); + } + } } } EmbedderSource::HuggingFace => { diff --git a/milli/src/vector/openai.rs b/milli/src/vector/openai.rs index 524f83b80..104decb66 100644 --- a/milli/src/vector/openai.rs +++ b/milli/src/vector/openai.rs @@ -17,6 +17,7 @@ pub struct Embedder { pub struct EmbedderOptions { pub api_key: Option, pub embedding_model: EmbeddingModel, + pub dimensions: Option, } #[derive( @@ -41,34 +42,50 @@ pub enum EmbeddingModel { #[serde(rename = "text-embedding-ada-002")] #[deserr(rename = "text-embedding-ada-002")] TextEmbeddingAda002, + + #[serde(rename = "text-embedding-3-small")] + #[deserr(rename = "text-embedding-3-small")] + TextEmbedding3Small, + + #[serde(rename = "text-embedding-3-large")] + #[deserr(rename = "text-embedding-3-large")] + TextEmbedding3Large, } impl EmbeddingModel { pub fn supported_models() -> &'static [&'static str] { - &["text-embedding-ada-002"] + &["text-embedding-ada-002", "text-embedding-3-small", "text-embedding-3-large"] } pub fn max_token(&self) -> usize { match self { EmbeddingModel::TextEmbeddingAda002 => 8191, + EmbeddingModel::TextEmbedding3Large => 8191, + EmbeddingModel::TextEmbedding3Small => 8191, } } - pub fn dimensions(&self) -> usize { + pub fn default_dimensions(&self) -> usize { match self { EmbeddingModel::TextEmbeddingAda002 => 1536, + EmbeddingModel::TextEmbedding3Large => 3072, + EmbeddingModel::TextEmbedding3Small => 1536, } } pub fn name(&self) -> &'static str { match self { EmbeddingModel::TextEmbeddingAda002 => "text-embedding-ada-002", + EmbeddingModel::TextEmbedding3Large => "text-embedding-3-large", + EmbeddingModel::TextEmbedding3Small => "text-embedding-3-small", } } pub fn from_name(name: &str) -> Option { match name { "text-embedding-ada-002" => Some(EmbeddingModel::TextEmbeddingAda002), + "text-embedding-3-large" => Some(EmbeddingModel::TextEmbedding3Large), + "text-embedding-3-small" => Some(EmbeddingModel::TextEmbedding3Small), _ => None, } } @@ -78,6 +95,20 @@ impl EmbeddingModel { EmbeddingModel::TextEmbeddingAda002 => { Some(DistributionShift { current_mean: 0.90, current_sigma: 0.08 }) } + EmbeddingModel::TextEmbedding3Large => { + Some(DistributionShift { current_mean: 0.70, current_sigma: 0.1 }) + } + EmbeddingModel::TextEmbedding3Small => { + Some(DistributionShift { current_mean: 0.75, current_sigma: 0.1 }) + } + } + } + + pub fn supports_overriding_dimensions(&self) -> bool { + match self { + EmbeddingModel::TextEmbeddingAda002 => false, + EmbeddingModel::TextEmbedding3Large => true, + EmbeddingModel::TextEmbedding3Small => true, } } } @@ -86,11 +117,11 @@ pub const OPENAI_EMBEDDINGS_URL: &str = "https://api.openai.com/v1/embeddings"; impl EmbedderOptions { pub fn with_default_model(api_key: Option) -> Self { - Self { api_key, embedding_model: Default::default() } + Self { api_key, embedding_model: Default::default(), dimensions: None } } pub fn with_embedding_model(api_key: Option, embedding_model: EmbeddingModel) -> Self { - Self { api_key, embedding_model } + Self { api_key, embedding_model, dimensions: None } } } @@ -237,7 +268,11 @@ impl Embedder { for text in texts { log::trace!("Received prompt: {}", text.as_ref()) } - let request = OpenAiRequest { model: self.options.embedding_model.name(), input: texts }; + let request = OpenAiRequest { + model: self.options.embedding_model.name(), + input: texts, + dimensions: self.overriden_dimensions(), + }; let response = client .post(OPENAI_EMBEDDINGS_URL) .json(&request) @@ -280,8 +315,7 @@ impl Embedder { } let mut tokens = encoded.as_slice(); - let mut embeddings_for_prompt = - Embeddings::new(self.options.embedding_model.dimensions()); + let mut embeddings_for_prompt = Embeddings::new(self.dimensions()); while tokens.len() > max_token_count { let window = &tokens[..max_token_count]; embeddings_for_prompt.push(self.embed_tokens(window, client).await?).unwrap(); @@ -322,8 +356,11 @@ impl Embedder { tokens: &[usize], client: &reqwest::Client, ) -> Result { - let request = - OpenAiTokensRequest { model: self.options.embedding_model.name(), input: tokens }; + let request = OpenAiTokensRequest { + model: self.options.embedding_model.name(), + input: tokens, + dimensions: self.overriden_dimensions(), + }; let response = client .post(OPENAI_EMBEDDINGS_URL) .json(&request) @@ -366,12 +403,24 @@ impl Embedder { } pub fn dimensions(&self) -> usize { - self.options.embedding_model.dimensions() + if self.options.embedding_model.supports_overriding_dimensions() { + self.options.dimensions.unwrap_or(self.options.embedding_model.default_dimensions()) + } else { + self.options.embedding_model.default_dimensions() + } } pub fn distribution(&self) -> Option { self.options.embedding_model.distribution() } + + fn overriden_dimensions(&self) -> Option { + if self.options.embedding_model.supports_overriding_dimensions() { + self.options.dimensions + } else { + None + } + } } // retrying in case of failure @@ -431,12 +480,16 @@ impl Retry { struct OpenAiRequest<'a, S: AsRef + serde::Serialize> { model: &'a str, input: &'a [S], + #[serde(skip_serializing_if = "Option::is_none")] + dimensions: Option, } #[derive(Debug, Serialize)] struct OpenAiTokensRequest<'a> { model: &'a str, input: &'a [usize], + #[serde(skip_serializing_if = "Option::is_none")] + dimensions: Option, } #[derive(Debug, Deserialize)] diff --git a/milli/src/vector/settings.rs b/milli/src/vector/settings.rs index 37fb80452..834a1c81d 100644 --- a/milli/src/vector/settings.rs +++ b/milli/src/vector/settings.rs @@ -1,6 +1,7 @@ use deserr::Deserr; use serde::{Deserialize, Serialize}; +use super::openai; use crate::prompt::PromptData; use crate::update::Setting; use crate::vector::EmbeddingConfig; @@ -82,7 +83,7 @@ impl EmbeddingSettings { Self::MODEL => &[EmbedderSource::HuggingFace, EmbedderSource::OpenAi], Self::REVISION => &[EmbedderSource::HuggingFace], Self::API_KEY => &[EmbedderSource::OpenAi], - Self::DIMENSIONS => &[EmbedderSource::UserProvided], + Self::DIMENSIONS => &[EmbedderSource::OpenAi, EmbedderSource::UserProvided], Self::DOCUMENT_TEMPLATE => &[EmbedderSource::HuggingFace, EmbedderSource::OpenAi], _other => unreachable!("unknown field"), } @@ -90,9 +91,13 @@ impl EmbeddingSettings { pub fn allowed_fields_for_source(source: EmbedderSource) -> &'static [&'static str] { match source { - EmbedderSource::OpenAi => { - &[Self::SOURCE, Self::MODEL, Self::API_KEY, Self::DOCUMENT_TEMPLATE] - } + EmbedderSource::OpenAi => &[ + Self::SOURCE, + Self::MODEL, + Self::API_KEY, + Self::DOCUMENT_TEMPLATE, + Self::DIMENSIONS, + ], EmbedderSource::HuggingFace => { &[Self::SOURCE, Self::MODEL, Self::REVISION, Self::DOCUMENT_TEMPLATE] } @@ -109,6 +114,17 @@ impl EmbeddingSettings { *source = Setting::Set(EmbedderSource::default()) } } + + pub(crate) fn apply_default_openai_model(setting: &mut Setting) { + if let Setting::Set(EmbeddingSettings { + source: Setting::Set(EmbedderSource::OpenAi), + model: model @ (Setting::NotSet | Setting::Reset), + .. + }) = setting + { + *model = Setting::Set(openai::EmbeddingModel::default().name().to_owned()) + } + } } #[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq, Deserr)] @@ -176,7 +192,7 @@ impl From for EmbeddingSettings { model: Setting::Set(options.embedding_model.name().to_owned()), revision: Setting::NotSet, api_key: options.api_key.map(Setting::Set).unwrap_or_default(), - dimensions: Setting::NotSet, + dimensions: options.dimensions.map(Setting::Set).unwrap_or_default(), document_template: Setting::Set(prompt.template), }, super::EmbedderOptions::UserProvided(options) => Self { @@ -208,6 +224,9 @@ impl From for EmbeddingConfig { if let Some(api_key) = api_key.set() { options.api_key = Some(api_key); } + if let Some(dimensions) = dimensions.set() { + options.dimensions = Some(dimensions); + } this.embedder_options = super::EmbedderOptions::OpenAi(options); } EmbedderSource::HuggingFace => {