From abc588daac6053ec2edbdcde3f5a2dc5eb7d50b8 Mon Sep 17 00:00:00 2001 From: sigoden Date: Fri, 21 Jun 2024 06:00:26 +0800 Subject: feat: support rerank (#620) --- src/config/input.rs | 21 ++++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) (limited to 'src/config/input.rs') diff --git a/src/config/input.rs b/src/config/input.rs index 892a41f..03ec999 100644 --- a/src/config/input.rs +++ b/src/config/input.rs @@ -169,20 +169,31 @@ impl Input { if !self.text.is_empty() { let rag = self.config.read().rag.clone(); if let Some(rag) = rag { - let (top_k, min_score_vector, min_score_text) = { + let (top_k, min_score_vector_search, min_score_fulltext_search) = { let config = self.config.read(); ( config.rag_top_k, - config.rag_min_score_vector, - config.rag_min_score_text, + config.rag_min_score_vector_search, + config.rag_min_score_fulltext_search, ) }; + let rerank = match self.config.read().rag_rerank_model.clone() { + Some(rerank_model_id) => { + let min_score = self.config.read().rag_min_score_rerank; + let rerank_model = + Model::retrieve_rerank(&self.config.read(), &rerank_model_id)?; + let rerank_client = init_client(&self.config, Some(rerank_model))?; + Some((rerank_client, min_score)) + } + None => None, + }; let embeddings = rag .search( &self.text, top_k, - min_score_vector, - min_score_text, + min_score_vector_search, + min_score_fulltext_search, + rerank, abort_signal, ) .await?; -- cgit v1.2.3