From 2eab71a641827e503b14952373aec82661192ba2 Mon Sep 17 00:00:00 2001 From: sigoden Date: Thu, 20 Jun 2024 11:26:45 +0800 Subject: feat: rag hybrid search (#618) --- src/rag/mod.rs | 101 +++++++++++++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 88 insertions(+), 13 deletions(-) (limited to 'src/rag/mod.rs') diff --git a/src/rag/mod.rs b/src/rag/mod.rs index 4ce280d..7939d98 100644 --- a/src/rag/mod.rs +++ b/src/rag/mod.rs @@ -1,3 +1,4 @@ +use self::bm25::*; use self::loader::*; use self::splitter::*; @@ -5,6 +6,7 @@ use crate::client::*; use crate::config::*; use crate::utils::*; +mod bm25; mod loader; mod splitter; @@ -16,8 +18,7 @@ use inquire::{required, validator::Validation, Select, Text}; use path_absolutize::Absolutize; use serde::{Deserialize, Serialize}; use serde_json::json; -use std::fmt::Debug; -use std::{io::BufReader, path::Path}; +use std::{collections::HashMap, fmt::Debug, io::BufReader, path::Path}; use tokio::sync::mpsc; pub struct Rag { @@ -26,6 +27,7 @@ pub struct Rag { path: String, model: Model, hnsw: Hnsw<'static, f32, DistCosine>, + bm25: BM25, data: RagData, } @@ -85,6 +87,7 @@ impl Rag { pub fn create(config: &GlobalConfig, name: &str, path: &Path, data: RagData) -> Result { let hnsw = data.build_hnsw(); + let bm25 = data.build_bm25(); let model = Model::retrieve_embedding(&config.read(), &data.model)?; let client = init_client(config, Some(model.clone()))?; let rag = Rag { @@ -94,6 +97,7 @@ impl Rag { data, model, hnsw, + bm25, }; Ok(rag) } @@ -194,12 +198,13 @@ impl Rag { &self, text: &str, top_k: usize, - minimum_score: f32, + min_score_vector: f32, + min_score_text: f32, abort_signal: AbortSignal, ) -> Result { let (stop_spinner_tx, _) = run_spinner("Searching").await; let ret = tokio::select! { - ret = self.search_impl(text, top_k, minimum_score) => { + ret = self.hybird_search(text, top_k, min_score_vector, min_score_text) => { ret } _ = watch_abort_signal(abort_signal) => { @@ -289,18 +294,44 @@ impl Rag { Ok(()) } - async fn search_impl( + async fn hybird_search( &self, - text: &str, + query: &str, top_k: usize, - minimum_score: f32, + min_score_vector: f32, + min_score_text: f32, ) -> Result> { + let (vector_search_result, text_search_result) = tokio::join!( + self.vector_search(query, top_k, min_score_vector), + self.text_search(query, top_k, min_score_text) + ); + let vector_search_ids = vector_search_result?; + let text_search_ids = text_search_result?; + let ids = reciprocal_rank_fusion(vector_search_ids, text_search_ids, 1.0, 1.0, top_k); + let output: Vec<_> = ids + .into_iter() + .filter_map(|id| { + let (file_index, document_index) = split_vector_id(id); + let file = self.data.files.get(file_index)?; + let document = file.documents.get(document_index)?; + Some(document.page_content.clone()) + }) + .collect(); + Ok(output) + } + + async fn vector_search( + &self, + query: &str, + top_k: usize, + min_score: f32, + ) -> Result> { let splitter = RecursiveCharacterTextSplitter::new( self.data.chunk_size, self.data.chunk_overlap, &DEFAULT_SEPARATES, ); - let texts = splitter.split_text(text); + let texts = splitter.split_text(query); let embeddings_data = EmbeddingsData::new(texts, true); let embeddings = self.create_embeddings(embeddings_data, None).await?; let output = self @@ -310,13 +341,10 @@ impl Rag { .flat_map(|list| { list.into_iter() .filter_map(|v| { - if v.distance < minimum_score { + if v.distance < min_score { return None; } - let (file_index, document_index) = split_vector_id(v.d_id); - let file = self.data.files.get(file_index)?; - let document = file.documents.get(document_index)?; - Some(document.page_content.clone()) + Some(v.d_id) }) .collect::>() }) @@ -324,6 +352,16 @@ impl Rag { Ok(output) } + async fn text_search( + &self, + query: &str, + top_k: usize, + min_score: f32, + ) -> Result> { + let output = self.bm25.search(query, top_k, Some(min_score as f64)); + Ok(output) + } + async fn create_embeddings( &self, data: EmbeddingsData, @@ -393,6 +431,17 @@ impl RagData { hnsw.parallel_insert(&list); hnsw } + + pub fn build_bm25(&self) -> BM25 { + let mut corpus = vec![]; + for (file_index, file) in self.files.iter().enumerate() { + for (document_index, document) in file.documents.iter().enumerate() { + let id = combine_vector_id(file_index, document_index); + corpus.push((id, document.page_content.clone())); + } + } + BM25::new(corpus, BM25Options::default()) + } } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -502,3 +551,29 @@ fn progress(spinner_message_tx: &Option>, message: let _ = tx.send(message); } } + +fn reciprocal_rank_fusion( + vector_search_ids: Vec, + text_search_ids: Vec, + vector_search_weight: f32, + text_search_weight: f32, + top_k: usize, +) -> Vec { + let rrf_k = top_k * 2; + let mut map: HashMap = HashMap::new(); + for (index, &item) in vector_search_ids.iter().enumerate() { + *map.entry(item).or_default() += + (1.0 / ((rrf_k + index + 1) as f32)) * vector_search_weight; + } + for (index, &item) in text_search_ids.iter().enumerate() { + *map.entry(item).or_default() += (1.0 / ((rrf_k + index + 1) as f32)) * text_search_weight; + } + let mut sorted_items: Vec<(VectorID, f32)> = map.into_iter().collect(); + sorted_items.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()); + + sorted_items + .into_iter() + .take(top_k) + .map(|(v, _)| v) + .collect() +} -- cgit v1.2.3