lc_rag/bm25/
chunked.rs

1// src/retrieval/bm25/chunked.rs
2//! BM25 Chunked Retriever - 支持 Parent-Child 文档结构的 BM25 检索器
3//!
4//! 基于 LlamaIndex AutoMerging 模式实现：
5//! - 文档拆分为 Parent + Leaf 两层
6//! - BM25 在 Leaf 层搜索
7//! - AutoMerging 合并同一 Parent 的多个 Leaf
8//! - 支持 Bincode 持久化
9
10use super::algorithm::{bm25_score, compute_idf, BM25Params};
11use super::tokenizer::Tokenizer;
12use lc_vector_stores::document_store::{ChunkDocument, ChunkedDocumentStoreTrait};
13use lc_vector_stores::{Document, VectorStoreError};
14use serde::{Deserialize, Serialize};
15use std::collections::HashMap;
16use std::path::Path;
17use std::sync::Arc;
18
19// ============================================================================
20// 数据结构定义
21// ============================================================================
22
23// ChunkDocument 现在在 document_store.rs 中定义，BM25 直接使用
24
25/// AutoMerging 配置
26#[derive(Debug, Clone, Serialize, Deserialize)]
27pub struct AutoMergingConfig {
28    /// 合并阈值：同一 Parent 下命中 Leaf 占比达到该比例时合并为 Parent 文档
29    pub merge_threshold: f32,
30    /// Leaf chunk 的大小（字符数）
31    pub leaf_chunk_size: usize,
32    /// Parent chunk 的大小（字符数）
33    pub parent_chunk_size: usize,
34    /// 每个 Parent 下期望的 Leaf 数量
35    pub leaves_per_parent: usize,
36}
37
38impl Default for AutoMergingConfig {
39    fn default() -> Self {
40        Self {
41            merge_threshold: 0.5,
42            leaf_chunk_size: 400,
43            parent_chunk_size: 2000,
44            leaves_per_parent: 5,
45        }
46    }
47}
48
49impl AutoMergingConfig {
50    /// 创建使用默认配置的 `AutoMergingConfig`
51    pub fn new() -> Self {
52        Self::default()
53    }
54
55    /// 设置合并阈值
56    pub fn with_threshold(mut self, threshold: f32) -> Self {
57        self.merge_threshold = threshold;
58        self
59    }
60
61    /// 设置 Leaf chunk 大小
62    pub fn with_leaf_size(mut self, size: usize) -> Self {
63        self.leaf_chunk_size = size;
64        self
65    }
66
67    /// 设置 Parent chunk 大小
68    pub fn with_parent_size(mut self, size: usize) -> Self {
69        self.parent_chunk_size = size;
70        self
71    }
72}
73
74/// AutoMerging 搜索结果
75#[derive(Debug, Clone)]
76pub struct ChunkedSearchResult {
77    /// 合并得到的 Parent 文档（若未触发合并则为 `None`）
78    pub merged_parent: Option<Document>,
79    /// 命中的 Leaf chunks
80    pub leaf_chunks: Vec<ChunkDocument>,
81    /// 该结果的 BM25 评分
82    pub score: f32,
83    /// 命中的查询词项
84    pub matched_terms: Vec<String>,
85    /// 所属 Parent 的 id
86    pub parent_id: String,
87}
88
89impl ChunkedSearchResult {
90    /// 返回合并结果的内容：优先返回 Parent 内容，否则拼接所有 Leaf 内容
91    pub fn content(&self) -> String {
92        if let Some(parent) = &self.merged_parent {
93            parent.content.clone()
94        } else {
95            self.leaf_chunks
96                .iter()
97                .map(|c| c.content.as_str())
98                .collect::<Vec<_>>()
99                .join("\n")
100        }
101    }
102
103    /// 是否触发了 AutoMerging 合并
104    pub fn is_merged(&self) -> bool {
105        self.merged_parent.is_some()
106    }
107}
108
109/// BM25 参数的可序列化版本
110#[derive(Debug, Clone, Serialize, Deserialize)]
111pub struct BM25ParamsData {
112    pub k1: f64,
113    pub b: f64,
114}
115
116impl From<BM25Params> for BM25ParamsData {
117    fn from(params: BM25Params) -> Self {
118        Self {
119            k1: params.k1,
120            b: params.b,
121        }
122    }
123}
124
125impl From<BM25ParamsData> for BM25Params {
126    fn from(data: BM25ParamsData) -> Self {
127        BM25Params::with_values(data.k1, data.b)
128    }
129}
130
131/// 可序列化的索引数据（不含内容，内容在ChunkedDocumentStore中）
132#[derive(Debug, Clone, Serialize, Deserialize)]
133pub struct ChunkedIndexData {
134    /// chunk 的 id 列表
135    pub chunk_id_list: Vec<String>,
136    /// 每个 chunk 的词频表
137    pub chunk_term_freqs: Vec<HashMap<String, usize>>,
138    /// 倒排索引：词项 -> (chunk 下标, 词频) 列表
139    pub term_index: HashMap<String, Vec<(usize, usize)>>,
140    /// Parent id -> 该 Parent 下 Leaf chunk 下标列表
141    pub parent_to_leaves: HashMap<String, Vec<usize>>,
142    /// 每个 chunk 的文档长度
143    pub doc_lengths: Vec<usize>,
144    /// 平均文档长度
145    pub avgdl: f64,
146    /// 文档数量
147    pub n_docs: usize,
148    /// BM25 参数
149    pub params: BM25ParamsData,
150    /// AutoMerging 配置
151    pub config: AutoMergingConfig,
152}
153
154// ============================================================================
155// ChunkedBM25Index 索引结构
156// ============================================================================
157
158/// 支持 Parent-Child 结构的 BM25 倒排索引
159pub struct ChunkedBM25Index<S: ChunkedDocumentStoreTrait = lc_vector_stores::ChunkedDocumentStore> {
160    store: Arc<S>,
161    chunk_id_list: Vec<String>,
162    chunk_term_freqs: Vec<HashMap<String, usize>>,
163    term_index: HashMap<String, Vec<(usize, usize)>>,
164    parent_to_leaves: HashMap<String, Vec<usize>>,
165    doc_lengths: Vec<usize>,
166    avgdl: f64,
167    n_docs: usize,
168    idf_cache: HashMap<String, f64>,
169    params: BM25Params,
170    tokenizer: Tokenizer,
171    config: AutoMergingConfig,
172}
173
174impl<S: ChunkedDocumentStoreTrait> ChunkedBM25Index<S> {
175    /// 使用默认配置创建索引
176    pub fn new(store: Arc<S>) -> Self {
177        Self::with_config(store, AutoMergingConfig::default())
178    }
179
180    /// 使用指定配置创建索引
181    pub fn with_config(store: Arc<S>, config: AutoMergingConfig) -> Self {
182        Self {
183            store,
184            chunk_id_list: Vec::new(),
185            chunk_term_freqs: Vec::new(),
186            term_index: HashMap::new(),
187            parent_to_leaves: HashMap::new(),
188            doc_lengths: Vec::new(),
189            avgdl: 0.0,
190            n_docs: 0,
191            idf_cache: HashMap::new(),
192            params: BM25Params::default(),
193            tokenizer: Tokenizer::new(),
194            config,
195        }
196    }
197
198    /// 使用指定 BM25 参数创建索引
199    pub fn with_params(store: Arc<S>, params: BM25Params) -> Self {
200        let mut index = Self::new(store);
201        index.params = params;
202        index
203    }
204
205    /// 添加chunk索引（内容已在store中）
206    pub fn add_chunk_index(
207        &mut self,
208        chunk_id: impl Into<String>,
209        parent_id: impl Into<String>,
210        content: &str,
211    ) {
212        let chunk_idx = self.n_docs;
213        let chunk_id = chunk_id.into();
214        let parent_id = parent_id.into();
215
216        let terms = self.tokenizer.tokenize(content);
217        let term_freq = self.compute_term_freq(&terms);
218
219        // 更新倒排索引
220        for (term, freq) in &term_freq {
221            self.term_index
222                .entry(term.clone())
223                .or_default()
224                .push((chunk_idx, *freq));
225        }
226
227        // 更新parent到chunk的映射
228        self.parent_to_leaves
229            .entry(parent_id)
230            .or_default()
231            .push(chunk_idx);
232
233        // 存储chunk_id和词频（BM25计算需要）
234        self.chunk_id_list.push(chunk_id);
235        self.chunk_term_freqs.push(term_freq.clone());
236
237        let doc_length: usize = term_freq.values().sum();
238        self.doc_lengths.push(doc_length);
239        self.n_docs += 1;
240        self.update_avgdl();
241        self.idf_cache.clear();
242    }
243
244    /// 批量添加chunk索引
245    pub fn add_chunk_indexes(&mut self, chunks: Vec<(String, String, String)>) {
246        for (chunk_id, parent_id, content) in chunks {
247            self.add_chunk_index(chunk_id, parent_id, &content);
248        }
249    }
250
251    fn compute_term_freq(&self, terms: &[String]) -> HashMap<String, usize> {
252        let mut freq = HashMap::new();
253        for term in terms {
254            *freq.entry(term.clone()).or_insert(0) += 1;
255        }
256        freq
257    }
258
259    fn update_avgdl(&mut self) {
260        if self.n_docs == 0 {
261            self.avgdl = 0.0;
262        } else {
263            let total: usize = self.doc_lengths.iter().sum();
264            self.avgdl = total as f64 / self.n_docs as f64;
265        }
266    }
267
268    fn compute_idf_for_term(&mut self, term: &str) -> f64 {
269        if let Some(idf) = self.idf_cache.get(term) {
270            return *idf;
271        }
272
273        let n = self.term_index.get(term).map(|v| v.len()).unwrap_or(0);
274        let idf = compute_idf(n, self.n_docs);
275        self.idf_cache.insert(term.to_string(), idf);
276        idf
277    }
278
279    /// 按 chunk 下标获取 chunk id
280    pub fn get_chunk_id(&self, chunk_idx: usize) -> Option<&String> {
281        self.chunk_id_list.get(chunk_idx)
282    }
283
284    /// 获取指定 Parent 下的所有 chunk id
285    pub fn get_chunk_ids_for_parent(&self, parent_id: &str) -> Vec<&String> {
286        self.parent_to_leaves
287            .get(parent_id)
288            .map(|indices| {
289                indices
290                    .iter()
291                    .filter_map(|idx| self.chunk_id_list.get(*idx))
292                    .collect()
293            })
294            .unwrap_or_default()
295    }
296
297    /// 返回 AutoMerging 配置
298    pub fn config(&self) -> &AutoMergingConfig {
299        &self.config
300    }
301
302    /// 返回已索引的文档数量
303    pub fn n_docs(&self) -> usize {
304        self.n_docs
305    }
306
307    /// 返回底层文档存储
308    pub fn store(&self) -> &Arc<S> {
309        &self.store
310    }
311
312    /// 清空索引数据
313    pub fn clear(&mut self) {
314        self.chunk_id_list.clear();
315        self.chunk_term_freqs.clear();
316        self.term_index.clear();
317        self.parent_to_leaves.clear();
318        self.doc_lengths.clear();
319        self.avgdl = 0.0;
320        self.n_docs = 0;
321        self.idf_cache.clear();
322    }
323}
324
325impl Default for ChunkedBM25Index<lc_vector_stores::ChunkedDocumentStore> {
326    fn default() -> Self {
327        Self::new(Arc::new(lc_vector_stores::ChunkedDocumentStore::new()))
328    }
329}
330
331// ============================================================================
332// ChunkedBM25Retriever 检索器
333// ============================================================================
334
335/// 基于 AutoMerging 的 BM25 检索器
336pub struct ChunkedBM25Retriever<
337    S: ChunkedDocumentStoreTrait = lc_vector_stores::ChunkedDocumentStore,
338> {
339    index: ChunkedBM25Index<S>,
340}
341
342impl<S: ChunkedDocumentStoreTrait> ChunkedBM25Retriever<S> {
343    /// 使用默认配置创建检索器
344    pub fn new(store: Arc<S>) -> Self {
345        Self {
346            index: ChunkedBM25Index::new(store),
347        }
348    }
349
350    /// 使用指定配置创建检索器
351    pub fn with_config(store: Arc<S>, config: AutoMergingConfig) -> Self {
352        Self {
353            index: ChunkedBM25Index::with_config(store, config),
354        }
355    }
356
357    /// 使用指定的 k1、b 参数创建检索器
358    pub fn with_params(store: Arc<S>, k1: f64, b: f64) -> Self {
359        Self {
360            index: ChunkedBM25Index::with_params(store, BM25Params::with_values(k1, b)),
361        }
362    }
363
364    /// 返回底层文档存储
365    pub fn store(&self) -> &Arc<S> {
366        self.index.store()
367    }
368
369    /// 添加单个 chunk 索引（内容已存储在 store 中）
370    pub fn add_chunk_index(
371        &mut self,
372        chunk_id: impl Into<String>,
373        parent_id: impl Into<String>,
374        content: &str,
375    ) {
376        self.index.add_chunk_index(chunk_id, parent_id, content);
377    }
378
379    /// 批量添加 chunk 索引
380    pub fn add_chunk_indexes(&mut self, chunks: Vec<(String, String, String)>) {
381        self.index.add_chunk_indexes(chunks);
382    }
383
384    /// 以同步方式添加文档：自动拆分 Parent/Leaf 并建立索引
385    pub fn add_document(&mut self, document: Document) -> Result<(), VectorStoreError> {
386        let parent_id = document
387            .id
388            .clone()
389            .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
390
391        // P0-1: 无 id 的文档先把预分配的 parent_id 挂到文档上再入库,
392        // 否则 store 内部会再生成一个新 uuid,导致 get_chunks_for_parent 用错 key 查空。
393        self.index.store.add_parent_document_blocking(
394            document.clone().with_id(parent_id.clone()),
395            self.index.config.leaf_chunk_size,
396        )?;
397
398        let chunks = self
399            .index
400            .store
401            .blocking_get_chunks_for_parent(&parent_id)?;
402
403        for chunk in chunks {
404            self.add_chunk_index(
405                chunk.chunk_id.clone(),
406                chunk.parent_id.clone(),
407                &chunk.content,
408            );
409        }
410
411        Ok(())
412    }
413
414    /// 以异步方式添加文档：自动拆分 Parent/Leaf 并建立索引
415    pub async fn add_document_async(&mut self, document: Document) -> Result<(), VectorStoreError> {
416        let parent_id = document
417            .id
418            .clone()
419            .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
420
421        self.index
422            .store
423            .add_parent_document(
424                document.clone().with_id(parent_id.clone()),
425                self.index.config.leaf_chunk_size,
426            )
427            .await?;
428
429        let chunks = self.index.store.get_chunks_for_parent(&parent_id).await?;
430
431        for chunk in chunks {
432            self.add_chunk_index(
433                chunk.chunk_id.clone(),
434                chunk.parent_id.clone(),
435                &chunk.content,
436            );
437        }
438
439        Ok(())
440    }
441
442    /// 批量以同步方式添加文档
443    pub fn add_documents(&mut self, documents: Vec<Document>) -> Result<(), VectorStoreError> {
444        for doc in documents {
445            self.add_document(doc)?;
446        }
447        Ok(())
448    }
449
450    /// 批量以异步方式添加文档
451    pub async fn add_documents_async(
452        &mut self,
453        documents: Vec<Document>,
454    ) -> Result<(), VectorStoreError> {
455        for doc in documents {
456            self.add_document_async(doc).await?;
457        }
458        Ok(())
459    }
460
461    /// 同步执行 BM25 检索，返回前 k 个 AutoMerging 结果
462    pub fn search(&mut self, query: &str, k: usize) -> Vec<ChunkedSearchResult> {
463        if self.index.n_docs == 0 {
464            return Vec::new();
465        }
466
467        let query_terms = self.index.tokenizer.tokenize(query);
468        if query_terms.is_empty() {
469            return Vec::new();
470        }
471
472        let idf_values: HashMap<String, f64> = query_terms
473            .iter()
474            .map(|t| (t.clone(), self.index.compute_idf_for_term(t)))
475            .collect();
476
477        let scored_chunks = self.score_chunks(&query_terms, &idf_values);
478
479        if scored_chunks.is_empty() {
480            return Vec::new();
481        }
482
483        let top_chunks: Vec<(usize, f64)> = scored_chunks.into_iter().take(k * 2).collect();
484
485        self.auto_merge_sync(top_chunks, k)
486    }
487
488    /// 异步执行 BM25 检索，返回前 k 个 AutoMerging 结果
489    pub async fn search_async(&mut self, query: &str, k: usize) -> Vec<ChunkedSearchResult> {
490        if self.index.n_docs == 0 {
491            return Vec::new();
492        }
493
494        let query_terms = self.index.tokenizer.tokenize(query);
495        if query_terms.is_empty() {
496            return Vec::new();
497        }
498
499        let idf_values: HashMap<String, f64> = query_terms
500            .iter()
501            .map(|t| (t.clone(), self.index.compute_idf_for_term(t)))
502            .collect();
503
504        let scored_chunks = self.score_chunks(&query_terms, &idf_values);
505
506        if scored_chunks.is_empty() {
507            return Vec::new();
508        }
509
510        let top_chunks: Vec<(usize, f64)> = scored_chunks.into_iter().take(k * 2).collect();
511
512        self.auto_merge_async(top_chunks, k).await
513    }
514
515    fn auto_merge_sync(
516        &self,
517        scored_chunks: Vec<(usize, f64)>,
518        k: usize,
519    ) -> Vec<ChunkedSearchResult> {
520        let threshold = self.index.config.merge_threshold;
521        let leaves_per_parent = self.index.config.leaves_per_parent;
522
523        let parent_stats = self.collect_parent_stats(&scored_chunks);
524
525        let mut results: Vec<ChunkedSearchResult> = Vec::new();
526
527        for (parent_id, matched_leaves) in parent_stats {
528            let ratio = matched_leaves.len() as f32 / leaves_per_parent as f32;
529
530            let avg_score =
531                matched_leaves.iter().map(|(_, s)| s).sum::<f64>() / matched_leaves.len() as f64;
532
533            let matched_terms = matched_leaves
534                .iter()
535                .filter_map(|(idx, _)| self.index.chunk_term_freqs.get(*idx))
536                .flat_map(|tf| tf.keys().cloned())
537                .collect::<Vec<_>>();
538
539            if ratio >= threshold {
540                let parent_doc = self
541                    .index
542                    .store()
543                    .get_parent_document_blocking(&parent_id)
544                    .ok()
545                    .flatten();
546
547                results.push(ChunkedSearchResult {
548                    merged_parent: parent_doc,
549                    leaf_chunks: Vec::new(),
550                    score: avg_score as f32,
551                    matched_terms,
552                    parent_id,
553                });
554            } else {
555                let leaf_chunks: Vec<ChunkDocument> = matched_leaves
556                    .iter()
557                    .filter_map(|(idx, _)| {
558                        let chunk_id = self.index.get_chunk_id(*idx)?;
559                        let chunk = self
560                            .index
561                            .store()
562                            .get_chunk_blocking(chunk_id)
563                            .ok()
564                            .flatten()?;
565                        Some(chunk)
566                    })
567                    .collect();
568
569                results.push(ChunkedSearchResult {
570                    merged_parent: None,
571                    leaf_chunks,
572                    score: avg_score as f32,
573                    matched_terms,
574                    parent_id,
575                });
576            }
577        }
578
579        results.sort_by(|a, b| {
580            b.score
581                .partial_cmp(&a.score)
582                .unwrap_or(std::cmp::Ordering::Equal)
583        });
584        results.into_iter().take(k).collect()
585    }
586
587    async fn auto_merge_async(
588        &self,
589        scored_chunks: Vec<(usize, f64)>,
590        k: usize,
591    ) -> Vec<ChunkedSearchResult> {
592        let threshold = self.index.config.merge_threshold;
593        let leaves_per_parent = self.index.config.leaves_per_parent;
594
595        let parent_stats = self.collect_parent_stats(&scored_chunks);
596
597        let mut results: Vec<ChunkedSearchResult> = Vec::new();
598
599        for (parent_id, matched_leaves) in parent_stats {
600            let ratio = matched_leaves.len() as f32 / leaves_per_parent as f32;
601
602            let avg_score =
603                matched_leaves.iter().map(|(_, s)| s).sum::<f64>() / matched_leaves.len() as f64;
604
605            let matched_terms = matched_leaves
606                .iter()
607                .filter_map(|(idx, _)| self.index.chunk_term_freqs.get(*idx))
608                .flat_map(|tf| tf.keys().cloned())
609                .collect::<Vec<_>>();
610
611            if ratio >= threshold {
612                let parent_doc = self
613                    .index
614                    .store()
615                    .get_parent_document(&parent_id)
616                    .await
617                    .ok()
618                    .flatten();
619
620                results.push(ChunkedSearchResult {
621                    merged_parent: parent_doc,
622                    leaf_chunks: Vec::new(),
623                    score: avg_score as f32,
624                    matched_terms,
625                    parent_id,
626                });
627            } else {
628                let mut leaf_chunks = Vec::new();
629                for (idx, _) in matched_leaves {
630                    if let Some(chunk_id) = self.index.get_chunk_id(idx) {
631                        match self.index.store().get_chunk(chunk_id).await {
632                            Ok(Some(chunk)) => leaf_chunks.push(chunk),
633                            Ok(None) => {}
634                            Err(e) => {
635                                // 不再静默吞错:读失败记日志,该 chunk 从结果中缺失
636                                log::error!(
637                                    "failed to read chunk `{}` during retrieval (chunk missing from results): {}",
638                                    chunk_id,
639                                    e
640                                );
641                            }
642                        }
643                    }
644                }
645
646                results.push(ChunkedSearchResult {
647                    merged_parent: None,
648                    leaf_chunks,
649                    score: avg_score as f32,
650                    matched_terms,
651                    parent_id,
652                });
653            }
654        }
655
656        results.sort_by(|a, b| {
657            b.score
658                .partial_cmp(&a.score)
659                .unwrap_or(std::cmp::Ordering::Equal)
660        });
661        results.into_iter().take(k).collect()
662    }
663
664    fn score_chunks(
665        &self,
666        query_terms: &[String],
667        idf_values: &HashMap<String, f64>,
668    ) -> Vec<(usize, f64)> {
669        let mut scored = Vec::new();
670
671        for chunk_idx in 0..self.index.n_docs {
672            if let Some(term_freqs) = self.index.chunk_term_freqs.get(chunk_idx) {
673                let doc_length = *self.index.doc_lengths.get(chunk_idx).unwrap_or(&0);
674
675                let score = bm25_score(
676                    query_terms,
677                    term_freqs,
678                    doc_length,
679                    self.index.avgdl,
680                    idf_values,
681                    &self.index.params,
682                );
683
684                if score > 0.0 {
685                    scored.push((chunk_idx, score));
686                }
687            }
688        }
689
690        scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
691        scored
692    }
693
694    fn collect_parent_stats(
695        &self,
696        scored_chunks: &[(usize, f64)],
697    ) -> HashMap<String, Vec<(usize, f64)>> {
698        let mut stats: HashMap<String, Vec<(usize, f64)>> = HashMap::new();
699
700        for (chunk_idx, score) in scored_chunks {
701            if let Some(chunk_id) = self.index.chunk_id_list.get(*chunk_idx) {
702                let parent_id = chunk_id.split("::").next().unwrap_or_default().to_string();
703                stats
704                    .entry(parent_id)
705                    .or_default()
706                    .push((*chunk_idx, *score));
707            }
708        }
709
710        stats
711    }
712
713    /// 按 Parent id 获取父文档
714    pub fn get_parent_document(&self, parent_id: &str) -> Option<Document> {
715        self.index
716            .store()
717            .get_parent_document_blocking(parent_id)
718            .ok()
719            .flatten()
720    }
721
722    /// 返回索引中的文档数量
723    pub fn len(&self) -> usize {
724        self.index.n_docs()
725    }
726
727    /// 索引是否为空
728    pub fn is_empty(&self) -> bool {
729        self.index.n_docs() == 0
730    }
731
732    /// 清空索引
733    pub fn clear(&mut self) {
734        self.index.clear();
735    }
736
737    /// 返回 AutoMerging 配置
738    pub fn config(&self) -> &AutoMergingConfig {
739        self.index.config()
740    }
741
742    // 持久化方法
743    /// 将索引数据序列化为 Bincode 并保存到指定路径
744    pub fn save(&self, path: impl AsRef<Path>) -> Result<(), Box<dyn std::error::Error>> {
745        let data = ChunkedIndexData {
746            chunk_id_list: self.index.chunk_id_list.clone(),
747            chunk_term_freqs: self.index.chunk_term_freqs.clone(),
748            term_index: self.index.term_index.clone(),
749            parent_to_leaves: self.index.parent_to_leaves.clone(),
750            doc_lengths: self.index.doc_lengths.clone(),
751            avgdl: self.index.avgdl,
752            n_docs: self.index.n_docs,
753            params: BM25ParamsData::from(self.index.params.clone()),
754            config: self.index.config.clone(),
755        };
756        let encoded = bincode::serialize(&data)?;
757        std::fs::write(path.as_ref(), encoded)?;
758        Ok(())
759    }
760}
761
762impl ChunkedBM25Retriever<lc_vector_stores::ChunkedDocumentStore> {
763    /// 从指定路径加载 Bincode 序列化的索引数据
764    pub fn load(
765        store: Arc<lc_vector_stores::ChunkedDocumentStore>,
766        path: impl AsRef<Path>,
767    ) -> Result<Self, Box<dyn std::error::Error>> {
768        let bytes = std::fs::read(path.as_ref())?;
769        let data: ChunkedIndexData = bincode::deserialize(&bytes)?;
770        let params: BM25Params = data.params.into();
771
772        Ok(Self {
773            index: ChunkedBM25Index {
774                store,
775                chunk_id_list: data.chunk_id_list,
776                chunk_term_freqs: data.chunk_term_freqs,
777                term_index: data.term_index,
778                parent_to_leaves: data.parent_to_leaves,
779                doc_lengths: data.doc_lengths,
780                avgdl: data.avgdl,
781                n_docs: data.n_docs,
782                idf_cache: HashMap::new(),
783                params,
784                tokenizer: Tokenizer::new(),
785                config: data.config,
786            },
787        })
788    }
789}
790
791impl Default for ChunkedBM25Retriever<lc_vector_stores::ChunkedDocumentStore> {
792    fn default() -> Self {
793        Self::new(Arc::new(lc_vector_stores::ChunkedDocumentStore::new()))
794    }
795}
lc_rag/bm25/chunked.rs

lc_rag/bm25/
chunked.rs
