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 => { 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_set(&model, EmbeddingSettings::MODEL, inferred_source, name)?;
check_unset(&revision, EmbeddingSettings::REVISION, 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 { impl Embedder {
/// Spawns a new embedder built from its options. /// Spawns a new embedder built from its options.
pub fn new(options: EmbedderOptions) -> std::result::Result<Self, NewEmbedderError> { pub fn new(options: EmbedderOptions) -> std::result::Result<Self, NewEmbedderError> {

View File

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

View File

@ -152,7 +152,18 @@ impl SettingsDiff {
ReindexAction::push_action(&mut reindex_action, ReindexAction::FullReindex); ReindexAction::push_action(&mut reindex_action, ReindexAction::FullReindex);
} }
if dimensions.apply(new_dimensions) { 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) { if url.apply(new_url) {
match source { match source {
@ -329,9 +340,12 @@ impl EmbeddingSettings {
Self::API_KEY => { Self::API_KEY => {
&[EmbedderSource::OpenAi, EmbedderSource::Ollama, EmbedderSource::Rest] &[EmbedderSource::OpenAi, EmbedderSource::Ollama, EmbedderSource::Rest]
} }
Self::DIMENSIONS => { Self::DIMENSIONS => &[
&[EmbedderSource::OpenAi, EmbedderSource::UserProvided, EmbedderSource::Rest] EmbedderSource::OpenAi,
} EmbedderSource::UserProvided,
EmbedderSource::Ollama,
EmbedderSource::Rest,
],
Self::DOCUMENT_TEMPLATE => &[ Self::DOCUMENT_TEMPLATE => &[
EmbedderSource::HuggingFace, EmbedderSource::HuggingFace,
EmbedderSource::OpenAi, EmbedderSource::OpenAi,
@ -377,6 +391,7 @@ impl EmbeddingSettings {
Self::DOCUMENT_TEMPLATE, Self::DOCUMENT_TEMPLATE,
Self::URL, Self::URL,
Self::API_KEY, Self::API_KEY,
Self::DIMENSIONS,
Self::DISTRIBUTION, Self::DISTRIBUTION,
], ],
EmbedderSource::UserProvided => &[Self::SOURCE, Self::DIMENSIONS, Self::DISTRIBUTION], EmbedderSource::UserProvided => &[Self::SOURCE, Self::DIMENSIONS, Self::DISTRIBUTION],
@ -486,12 +501,13 @@ impl From<EmbeddingConfig> for EmbeddingSettings {
url, url,
api_key, api_key,
distribution, distribution,
dimensions,
}) => Self { }) => Self {
source: Setting::Set(EmbedderSource::Ollama), source: Setting::Set(EmbedderSource::Ollama),
model: Setting::Set(embedding_model), model: Setting::Set(embedding_model),
revision: Setting::NotSet, revision: Setting::NotSet,
api_key: api_key.map(Setting::Set).unwrap_or_default(), 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), document_template: Setting::Set(prompt.template),
url: url.map(Setting::Set).unwrap_or_default(), url: url.map(Setting::Set).unwrap_or_default(),
request: Setting::NotSet, request: Setting::NotSet,
@ -583,6 +599,7 @@ impl From<EmbeddingSettings> for EmbeddingConfig {
super::ollama::EmbedderOptions::with_default_model( super::ollama::EmbedderOptions::with_default_model(
api_key.set(), api_key.set(),
url.set(), url.set(),
dimensions.set(),
); );
if let Some(model) = model.set() { if let Some(model) = model.set() {
options.embedding_model = model; options.embedding_model = model;