Add ranking_score_threshold to milli

This commit is contained in:
Louis Dureuil
2024-04-11 19:04:06 +02:00
parent 75d5c0ae1f
commit aac1d769a7
5 changed files with 46 additions and 0 deletions

View File

@ -50,6 +50,7 @@ pub struct Search<'a> {
index: &'a Index,
semantic: Option<SemanticSearch>,
time_budget: TimeBudget,
ranking_score_threshold: Option<f64>,
}
impl<'a> Search<'a> {
@ -70,6 +71,7 @@ impl<'a> Search<'a> {
index,
semantic: None,
time_budget: TimeBudget::max(),
ranking_score_threshold: None,
}
}
@ -146,6 +148,14 @@ impl<'a> Search<'a> {
self
}
pub fn ranking_score_threshold(
&mut self,
ranking_score_threshold: Option<f64>,
) -> &mut Search<'a> {
self.ranking_score_threshold = ranking_score_threshold;
self
}
pub fn execute_for_candidates(&self, has_vector_search: bool) -> Result<RoaringBitmap> {
if has_vector_search {
let ctx = SearchContext::new(self.index, self.rtxn)?;
@ -184,6 +194,7 @@ impl<'a> Search<'a> {
embedder_name,
embedder,
self.time_budget.clone(),
self.ranking_score_threshold,
)?
}
_ => execute_search(
@ -201,6 +212,7 @@ impl<'a> Search<'a> {
&mut DefaultSearchLogger,
&mut DefaultSearchLogger,
self.time_budget.clone(),
self.ranking_score_threshold,
)?,
};
@ -239,6 +251,7 @@ impl fmt::Debug for Search<'_> {
index: _,
semantic,
time_budget,
ranking_score_threshold,
} = self;
f.debug_struct("Search")
.field("query", query)
@ -257,6 +270,7 @@ impl fmt::Debug for Search<'_> {
&semantic.as_ref().map(|semantic| &semantic.embedder_name),
)
.field("time_budget", time_budget)
.field("ranking_score_threshold", ranking_score_threshold)
.finish()
}
}