Lazily embed, don't fail hybrid search on embedding failure

milli/src/index.rs

@@ -1499,14 +1499,6 @@ impl Index {
             .unwrap_or_default())
     }
 
-    pub fn default_embedding_name(&self, rtxn: &RoTxn<'_>) -> Result<String> {
-        let configs = self.embedding_configs(rtxn)?;
-        Ok(match configs.as_slice() {
-            [(ref first_name, _)] => first_name.clone(),
-            _ => "default".to_owned(),
-        })
-    }
-
     pub(crate) fn put_search_cutoff(&self, wtxn: &mut RwTxn<'_>, cutoff: u64) -> heed::Result<()> {
         self.main.remap_types::<Str, BEU64>().put(wtxn, main_key::SEARCH_CUTOFF, &cutoff)
     }

milli/src/lib.rs

@@ -61,7 +61,7 @@ pub use self::index::Index;
 pub use self::search::facet::{FacetValueHit, SearchForFacetValues};
 pub use self::search::{
     FacetDistribution, Filter, FormatOptions, MatchBounds, MatcherBuilder, MatchingWords, OrderBy,
-    Search, SearchResult, TermsMatchingStrategy, DEFAULT_VALUES_PER_FACET,
+    Search, SearchResult, SemanticSearch, TermsMatchingStrategy, DEFAULT_VALUES_PER_FACET,
 };
 
 pub type Result<T> = std::result::Result<T, error::Error>;

milli/src/search/facet/search.rs

@@ -92,9 +92,15 @@ impl<'a> SearchForFacetValues<'a> {
             None => return Ok(Vec::new()),
         };
 
-        let search_candidates = self
-            .search_query
-            .execute_for_candidates(self.is_hybrid || self.search_query.vector.is_some())?;
+        let search_candidates = self.search_query.execute_for_candidates(
+            self.is_hybrid
+                || self
+                    .search_query
+                    .semantic
+                    .as_ref()
+                    .and_then(|semantic| semantic.vector.as_ref())
+                    .is_some(),
+        )?;
 
         let mut results = match index.sort_facet_values_by(rtxn)?.get(&self.facet) {
             OrderBy::Lexicographic => ValuesCollection::by_lexicographic(self.max_values),

milli/src/search/hybrid.rs

@@ -4,6 +4,7 @@ use itertools::Itertools;
 use roaring::RoaringBitmap;
 
 use crate::score_details::{ScoreDetails, ScoreValue, ScoringStrategy};
+use crate::search::SemanticSearch;
 use crate::{MatchingWords, Result, Search, SearchResult};
 
 struct ScoreWithRatioResult {
@@ -126,7 +127,6 @@ impl<'a> Search<'a> {
         // create separate keyword and semantic searches
         let mut search = Search {
             query: self.query.clone(),
-            vector: self.vector.clone(),
             filter: self.filter.clone(),
             offset: 0,
             limit: self.limit + self.offset,
@@ -139,26 +139,41 @@ impl<'a> Search<'a> {
             exhaustive_number_hits: self.exhaustive_number_hits,
             rtxn: self.rtxn,
             index: self.index,
-            distribution_shift: self.distribution_shift,
-            embedder_name: self.embedder_name.clone(),
+            semantic: self.semantic.clone(),
             time_budget: self.time_budget.clone(),
         };
 
-        let vector_query = search.vector.take();
+        let semantic = search.semantic.take();
         let keyword_results = search.execute()?;
 
-        // skip semantic search if we don't have a vector query (placeholder search)
-        let Some(vector_query) = vector_query else {
-            return Ok(keyword_results);
-        };
-
         // completely skip semantic search if the results of the keyword search are good enough
         if self.results_good_enough(&keyword_results, semantic_ratio) {
             return Ok(keyword_results);
         }
 
-        search.vector = Some(vector_query);
-        search.query = None;
+        // no vector search against placeholder search
+        let Some(query) = search.query.take() else { return Ok(keyword_results) };
+        // no embedder, no semantic search
+        let Some(SemanticSearch { vector, embedder_name, embedder }) = semantic else {
+            return Ok(keyword_results);
+        };
+
+        let vector_query = match vector {
+            Some(vector_query) => vector_query,
+            None => {
+                // attempt to embed the vector
+                match embedder.embed_one(query) {
+                    Ok(embedding) => embedding,
+                    Err(error) => {
+                        tracing::error!(error=%error, "Embedding failed");
+                        return Ok(keyword_results);
+                    }
+                }
+            }
+        };
+
+        search.semantic =
+            Some(SemanticSearch { vector: Some(vector_query), embedder_name, embedder });
 
         // TODO: would be better to have two distinct functions at this point
         let vector_results = search.execute()?;

milli/src/search/mod.rs

@@ -1,4 +1,5 @@
 use std::fmt;
+use std::sync::Arc;
 
 use levenshtein_automata::{LevenshteinAutomatonBuilder as LevBuilder, DFA};
 use once_cell::sync::Lazy;
@@ -8,7 +9,7 @@ pub use self::facet::{FacetDistribution, Filter, OrderBy, DEFAULT_VALUES_PER_FAC
 pub use self::new::matches::{FormatOptions, MatchBounds, MatcherBuilder, MatchingWords};
 use self::new::{execute_vector_search, PartialSearchResult};
 use crate::score_details::{ScoreDetails, ScoringStrategy};
-use crate::vector::DistributionShift;
+use crate::vector::Embedder;
 use crate::{
     execute_search, filtered_universe, AscDesc, DefaultSearchLogger, DocumentId, Index, Result,
     SearchContext, TimeBudget,
@@ -24,9 +25,15 @@ mod fst_utils;
 pub mod hybrid;
 pub mod new;
 
+#[derive(Debug, Clone)]
+pub struct SemanticSearch {
+    vector: Option<Vec<f32>>,
+    embedder_name: String,
+    embedder: Arc<Embedder>,
+}
+
 pub struct Search<'a> {
     query: Option<String>,
-    vector: Option<Vec<f32>>,
     // this should be linked to the String in the query
     filter: Option<Filter<'a>>,
     offset: usize,
@@ -38,12 +45,9 @@ pub struct Search<'a> {
     scoring_strategy: ScoringStrategy,
     words_limit: usize,
     exhaustive_number_hits: bool,
-    /// TODO: Add semantic ratio or pass it directly to execute_hybrid()
     rtxn: &'a heed::RoTxn<'a>,
     index: &'a Index,
-    distribution_shift: Option<DistributionShift>,
-    embedder_name: Option<String>,
-
+    semantic: Option<SemanticSearch>,
     time_budget: TimeBudget,
 }
 
@@ -51,7 +55,6 @@ impl<'a> Search<'a> {
     pub fn new(rtxn: &'a heed::RoTxn, index: &'a Index) -> Search<'a> {
         Search {
             query: None,
-            vector: None,
             filter: None,
             offset: 0,
             limit: 20,
@@ -64,8 +67,7 @@ impl<'a> Search<'a> {
             words_limit: 10,
             rtxn,
             index,
-            distribution_shift: None,
-            embedder_name: None,
+            semantic: None,
             time_budget: TimeBudget::max(),
         }
     }
@@ -75,8 +77,13 @@ impl<'a> Search<'a> {
         self
     }
 
-    pub fn vector(&mut self, vector: Vec<f32>) -> &mut Search<'a> {
-        self.vector = Some(vector);
+    pub fn semantic(
+        &mut self,
+        embedder_name: String,
+        embedder: Arc<Embedder>,
+        vector: Option<Vec<f32>>,
+    ) -> &mut Search<'a> {
+        self.semantic = Some(SemanticSearch { embedder_name, embedder, vector });
         self
     }
 
@@ -133,19 +140,6 @@ impl<'a> Search<'a> {
         self
     }
 
-    pub fn distribution_shift(
-        &mut self,
-        distribution_shift: Option<DistributionShift>,
-    ) -> &mut Search<'a> {
-        self.distribution_shift = distribution_shift;
-        self
-    }
-
-    pub fn embedder_name(&mut self, embedder_name: impl Into<String>) -> &mut Search<'a> {
-        self.embedder_name = Some(embedder_name.into());
-        self
-    }
-
     pub fn time_budget(&mut self, time_budget: TimeBudget) -> &mut Search<'a> {
         self.time_budget = time_budget;
         self
@@ -161,15 +155,6 @@ impl<'a> Search<'a> {
     }
 
     pub fn execute(&self) -> Result<SearchResult> {
-        let embedder_name;
-        let embedder_name = match &self.embedder_name {
-            Some(embedder_name) => embedder_name,
-            None => {
-                embedder_name = self.index.default_embedding_name(self.rtxn)?;
-                &embedder_name
-            }
-        };
-
         let mut ctx = SearchContext::new(self.index, self.rtxn);
 
         if let Some(searchable_attributes) = self.searchable_attributes {
@@ -184,21 +169,23 @@ impl<'a> Search<'a> {
             document_scores,
             degraded,
             used_negative_operator,
-        } = match self.vector.as_ref() {
-            Some(vector) => execute_vector_search(
-                &mut ctx,
-                vector,
-                self.scoring_strategy,
-                universe,
-                &self.sort_criteria,
-                self.geo_strategy,
-                self.offset,
-                self.limit,
-                self.distribution_shift,
-                embedder_name,
-                self.time_budget.clone(),
-            )?,
-            None => execute_search(
+        } = match self.semantic.as_ref() {
+            Some(SemanticSearch { vector: Some(vector), embedder_name, embedder }) => {
+                execute_vector_search(
+                    &mut ctx,
+                    vector,
+                    self.scoring_strategy,
+                    universe,
+                    &self.sort_criteria,
+                    self.geo_strategy,
+                    self.offset,
+                    self.limit,
+                    embedder_name,
+                    embedder,
+                    self.time_budget.clone(),
+                )?
+            }
+            _ => execute_search(
                 &mut ctx,
                 self.query.as_deref(),
                 self.terms_matching_strategy,
@@ -237,7 +224,6 @@ impl fmt::Debug for Search<'_> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         let Search {
             query,
-            vector: _,
             filter,
             offset,
             limit,
@@ -250,8 +236,7 @@ impl fmt::Debug for Search<'_> {
             exhaustive_number_hits,
             rtxn: _,
             index: _,
-            distribution_shift,
-            embedder_name,
+            semantic,
             time_budget,
         } = self;
         f.debug_struct("Search")
@@ -266,8 +251,10 @@ impl fmt::Debug for Search<'_> {
             .field("scoring_strategy", scoring_strategy)
             .field("exhaustive_number_hits", exhaustive_number_hits)
             .field("words_limit", words_limit)
-            .field("distribution_shift", distribution_shift)
-            .field("embedder_name", embedder_name)
+            .field(
+                "semantic.embedder_name",
+                &semantic.as_ref().map(|semantic| &semantic.embedder_name),
+            )
             .field("time_budget", time_budget)
             .finish()
     }

milli/src/search/new/mod.rs

@@ -52,7 +52,7 @@ use self::vector_sort::VectorSort;
 use crate::error::FieldIdMapMissingEntry;
 use crate::score_details::{ScoreDetails, ScoringStrategy};
 use crate::search::new::distinct::apply_distinct_rule;
-use crate::vector::DistributionShift;
+use crate::vector::Embedder;
 use crate::{
     AscDesc, DocumentId, FieldId, Filter, Index, Member, Result, TermsMatchingStrategy, TimeBudget,
     UserError,
@@ -298,8 +298,8 @@ fn get_ranking_rules_for_vector<'ctx>(
     geo_strategy: geo_sort::Strategy,
     limit_plus_offset: usize,
     target: &[f32],
-    distribution_shift: Option<DistributionShift>,
     embedder_name: &str,
+    embedder: &Embedder,
 ) -> Result<Vec<BoxRankingRule<'ctx, PlaceholderQuery>>> {
     // query graph search
 
@@ -325,8 +325,8 @@ fn get_ranking_rules_for_vector<'ctx>(
                         target.to_vec(),
                         vector_candidates,
                         limit_plus_offset,
-                        distribution_shift,
                         embedder_name,
+                        embedder,
                     )?;
                     ranking_rules.push(Box::new(vector_sort));
                     vector = true;
@@ -548,8 +548,8 @@ pub fn execute_vector_search(
     geo_strategy: geo_sort::Strategy,
     from: usize,
     length: usize,
-    distribution_shift: Option<DistributionShift>,
     embedder_name: &str,
+    embedder: &Embedder,
     time_budget: TimeBudget,
 ) -> Result<PartialSearchResult> {
     check_sort_criteria(ctx, sort_criteria.as_ref())?;
@@ -562,8 +562,8 @@ pub fn execute_vector_search(
         geo_strategy,
         from + length,
         vector,
-        distribution_shift,
         embedder_name,
+        embedder,
     )?;
 
     let mut placeholder_search_logger = logger::DefaultSearchLogger;

milli/src/search/new/vector_sort.rs

@@ -5,7 +5,7 @@ use roaring::RoaringBitmap;
 
 use super::ranking_rules::{RankingRule, RankingRuleOutput, RankingRuleQueryTrait};
 use crate::score_details::{self, ScoreDetails};
-use crate::vector::DistributionShift;
+use crate::vector::{DistributionShift, Embedder};
 use crate::{DocumentId, Result, SearchContext, SearchLogger};
 
 pub struct VectorSort<Q: RankingRuleQueryTrait> {
@@ -24,8 +24,8 @@ impl<Q: RankingRuleQueryTrait> VectorSort<Q> {
         target: Vec<f32>,
         vector_candidates: RoaringBitmap,
         limit: usize,
-        distribution_shift: Option<DistributionShift>,
         embedder_name: &str,
+        embedder: &Embedder,
     ) -> Result<Self> {
         let embedder_index = ctx
             .index
@@ -39,7 +39,7 @@ impl<Q: RankingRuleQueryTrait> VectorSort<Q> {
             vector_candidates,
             cached_sorted_docids: Default::default(),
             limit,
-            distribution_shift,
+            distribution_shift: embedder.distribution(),
             embedder_index,
         })
     }