diff options
| author | sigoden <sigoden@gmail.com> | 2024-06-20 11:26:45 +0800 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2024-06-20 11:26:45 +0800 |
| commit | 2eab71a641827e503b14952373aec82661192ba2 (patch) | |
| tree | b2f1fd1d8b988391d85a8058dcaaadd6ef379533 /src/config | |
| parent | 3b3d39cef0211b607d51360d0739022f953ba1a7 (diff) | |
| download | aichat-2eab71a641827e503b14952373aec82661192ba2.tar.gz | |
feat: rag hybrid search (#618)
Diffstat (limited to 'src/config')
| -rw-r--r-- | src/config/input.rs | 16 | ||||
| -rw-r--r-- | src/config/mod.rs | 13 |
2 files changed, 17 insertions, 12 deletions
diff --git a/src/config/input.rs b/src/config/input.rs index 935d517..892a41f 100644 --- a/src/config/input.rs +++ b/src/config/input.rs @@ -169,12 +169,22 @@ impl Input { if !self.text.is_empty() { let rag = self.config.read().rag.clone(); if let Some(rag) = rag { - let (top_k, minimum_score) = { + let (top_k, min_score_vector, min_score_text) = { let config = self.config.read(); - (config.rag_top_k, config.rag_minimum_score) + ( + config.rag_top_k, + config.rag_min_score_vector, + config.rag_min_score_text, + ) }; let embeddings = rag - .search(&self.text, top_k, minimum_score, abort_signal) + .search( + &self.text, + top_k, + min_score_vector, + min_score_text, + abort_signal, + ) .await?; let text = self.config.read().rag_template(&embeddings, &self.text); self.patched_text = Some(text); diff --git a/src/config/mod.rs b/src/config/mod.rs index b2e4517..509e4de 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -105,7 +105,8 @@ pub struct Config { pub rag_chunk_size: Option<usize>, pub rag_chunk_overlap: Option<usize>, pub rag_top_k: usize, - pub rag_minimum_score: f32, + pub rag_min_score_vector: f32, + pub rag_min_score_text: f32, pub rag_template: Option<String>, pub compress_threshold: usize, pub summarize_prompt: Option<String>, @@ -158,7 +159,8 @@ impl Default for Config { rag_chunk_size: None, rag_chunk_overlap: None, rag_top_k: 4, - rag_minimum_score: 0.0, + rag_min_score_vector: 0.0, + rag_min_score_text: 0.0, rag_template: None, compress_threshold: 4000, summarize_prompt: None, @@ -441,7 +443,6 @@ impl Config { ("temperature", format_option_value(&role.temperature())), ("top_p", format_option_value(&role.top_p())), ("rag_top_k", self.rag_top_k.to_string()), - ("rag_minimum_score", self.rag_minimum_score.to_string()), ("function_calling", self.function_calling.to_string()), ("compress_threshold", self.compress_threshold.to_string()), ("dry_run", self.dry_run.to_string()), @@ -494,11 +495,6 @@ impl Config { self.rag_top_k = value; } } - "rag_minimum_score" => { - if let Some(value) = parse_value(value)? { - self.rag_minimum_score = value; - } - } "function_calling" => { let value = value.parse().with_context(|| "Invalid value")?; self.function_calling = value; @@ -1057,7 +1053,6 @@ impl Config { "temperature", "top_p", "rag_top_k", - "rag_minimum_score", "function_calling", "compress_threshold", "save", |