使用 Java 实现混合检索策略（BM25+向量）提升 RAG 召回精准度与复杂业务匹配能力
大家好，今天我们将深入探讨如何使用 Java 实现混合检索策略，特别是结合 BM25 和向量搜索，来提升 RAG（Retrieval-Augmented Generation，检索增强生成）系统的召回精准度，并使其更好地适应复杂业务场景。RAG 是一种将检索和生成模型结合起来的技术，它首先从知识库中检索相关文档，然后利用这些文档来指导生成模型生成更准确、更相关的答案。检索环节的质量直接决定了 RAG 系统的性能，因此优化检索策略至关重要。
RAG 系统中的检索挑战
传统的检索方法，如基于关键词匹配的 BM25，在处理精确匹配和常见查询时表现良好，但面对语义相似性、上下文理解以及复杂的业务逻辑时，往往力不从心。例如，用户可能使用不同的措辞来表达相同的含义，或者查询涉及多个实体和关系，这些都超出了关键词匹配的能力范围。
向量搜索，特别是基于嵌入（embedding）的搜索，通过将文本转换为高维向量空间中的点，可以捕捉语义相似性。然而，单纯的向量搜索有时会忽略关键词的重要性，导致召回结果与用户的意图存在偏差。
因此，我们需要一种混合检索策略，既能利用 BM25 的精确匹配能力，又能利用向量搜索的语义理解能力，从而实现更精准、更全面的召回。
BM25 算法原理及 Java 实现
BM25（Best Matching 25）是一种基于概率检索模型的排序函数，用于评估文档与查询之间的相关性。它考虑了词频、文档长度以及文档频率等因素。
BM25 公式:
score(D, Q) = Σ IDF(qi) * ((f(qi, D) * (k1 + 1)) / (f(qi, D) + k1 * (1 - b + b * (|D| / avgdl))))其中：
D：文档Q：查询qi：查询中的第 i 个词f(qi, D)：词qi在文档D中的词频|D|：文档D的长度（词数）avgdl：所有文档的平均长度IDF(qi)：词qi的逆文档频率k1和b：可调节的参数，通常 k1 在 1.2 到 2.0 之间，b 在 0.75 左右。
IDF 公式:
IDF(qi) = log((N - n(qi) + 0.5) / (n(qi) + 0.5))其中：
N：文档总数n(qi)：包含词qi的文档数
Java 代码实现:
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
public class BM25 {
private List<String> documents;
private double avgdl;
private Map<String, Integer> documentFrequencies = new HashMap<>();
private int N;
private double k1 = 1.2;
private double b = 0.75;
public BM25(List<String> documents) {
this.documents = documents;
this.N = documents.size();
this.avgdl = documents.stream().mapToInt(doc -> doc.split(" ").length).average().orElse(0.0);
calculateDocumentFrequencies();
}
private void calculateDocumentFrequencies() {
for (String document : documents) {
Arrays.stream(document.split(" ")).forEach(term -> {
documentFrequencies.put(term, documentFrequencies.getOrDefault(term, 0) + 1);
});
}
}
private double calculateIDF(String term) {
int nqi = documentFrequencies.getOrDefault(term, 0);
return Math.log((N - nqi + 0.5) / (nqi + 0.5));
}
public double score(String document, String query) {
double score = 0.0;
String[] queryTerms = query.split(" ");
int documentLength = document.split(" ").length;
for (String term : queryTerms) {
double idf = calculateIDF(term);
int termFrequency = 0;
for(String word : document.split(" ")){
if(word.equals(term)){
termFrequency++;
}
}
score += idf * ((termFrequency * (k1 + 1)) / (termFrequency + k1 * (1 - b + b * (documentLength / avgdl))));
}
return score;
}
public static void main(String[] args) {
List<String> documents = Arrays.asList(
"This is the first document about information retrieval.",
"This is the second document.",
"And this is the third one.",
"Is this the first document?"
);
BM25 bm25 = new BM25(documents);
String query = "first document";
for (int i = 0; i < documents.size(); i++) {
double score = bm25.score(documents.get(i), query);
System.out.println("Document " + (i + 1) + " score: " + score);
}
}
}这段代码实现了一个简单的 BM25 评分器。它首先计算每个词的逆文档频率 (IDF)，然后根据 BM25 公式计算文档与查询之间的相关性得分。 k1 和 b 是可调参数，需要根据具体数据集进行调整以获得最佳性能。
向量搜索原理及 Java 实现
向量搜索通过将文本转换为向量表示，然后在向量空间中查找与查询向量最相似的文档向量来实现。常用的向量化方法包括 TF-IDF、Word2Vec、GloVe 和 Transformer 模型（如 BERT、Sentence-BERT）。
流程:
- 文本向量化: 使用预训练模型或自定义模型将文本转换为向量。
- 索引构建: 将文档向量存储在向量数据库中，例如 Faiss、Annoy 或 Milvus。
- 查询向量化: 将查询文本转换为向量。
- 相似度计算: 计算查询向量与所有文档向量之间的相似度，常用的相似度度量包括余弦相似度、欧氏距离和点积。
- 排序和召回: 根据相似度得分对文档进行排序，并召回得分最高的文档。
Java 代码实现 (使用 Sentence-BERT 和 cosine similarity):
为了实现向量搜索，我们需要一个现成的 Sentence-BERT 模型和一个向量数据库。这里我们简化实现，使用简单的余弦相似度计算，并假设我们已经有了文档的向量表示。实际应用中，你需要集成 Sentence-BERT 模型（可以使用 Java 的 TensorFlow 或 PyTorch 桥接库）和向量数据库。
import java.util.Arrays;
import java.util.List;
public class VectorSearch {
// 假设我们已经有了文档的向量表示
private static final List<double[]> documentVectors = Arrays.asList(
new double[]{0.1, 0.2, 0.3, 0.4},
new double[]{0.5, 0.6, 0.7, 0.8},
new double[]{0.9, 0.1, 0.2, 0.3},
new double[]{0.4, 0.5, 0.6, 0.7}
);
public static double cosineSimilarity(double[] vectorA, double[] vectorB) {
double dotProduct = 0.0;
double magnitudeA = 0.0;
double magnitudeB = 0.0;
for (int i = 0; i < vectorA.length; i++) {
dotProduct += vectorA[i] * vectorB[i];
magnitudeA += Math.pow(vectorA[i], 2);
magnitudeB += Math.pow(vectorB[i], 2);
}
magnitudeA = Math.sqrt(magnitudeA);
magnitudeB = Math.sqrt(magnitudeB);
return dotProduct / (magnitudeA * magnitudeB);
}
public static int search(double[] queryVector) {
int bestMatchIndex = -1;
double maxSimilarity = -1.0;
for (int i = 0; i < documentVectors.size(); i++) {
double similarity = cosineSimilarity(queryVector, documentVectors.get(i));
if (similarity > maxSimilarity) {
maxSimilarity = similarity;
bestMatchIndex = i;
}
}
return bestMatchIndex;
}
public static void main(String[] args) {
// 假设查询向量
double[] queryVector = {0.2, 0.3, 0.4, 0.5};
int bestMatchIndex = search(queryVector);
System.out.println("Best match document index: " + bestMatchIndex);
System.out.println("Cosine Similarity: " + cosineSimilarity(queryVector, documentVectors.get(bestMatchIndex)));
}
}这段代码演示了如何使用余弦相似度进行向量搜索。在实际应用中，你需要使用更复杂的向量数据库和更强大的向量化模型。 Sentence-BERT 是一个不错的选择，因为它专门用于生成句子级别的嵌入，能够更好地捕捉语义信息。
混合检索策略：BM25 + 向量搜索
为了结合 BM25 和向量搜索的优势，我们可以采用以下几种混合策略：
-
Rank Fusion (排序融合): 分别使用 BM25 和向量搜索对文档进行排序，然后将两个排序结果进行融合。常用的融合方法包括：
- 线性加权: 为 BM25 和向量搜索的得分分配不同的权重，然后将加权后的得分相加。
- Reciprocal Rank Fusion (RRF): 根据文档在每个排序列表中的排名来计算融合得分。
-
Filtering (过滤): 先使用 BM25 或向量搜索进行初步筛选，然后使用另一种方法对筛选结果进行重新排序。
-
Two-Stage Retrieval (两阶段检索): 先用BM25召回一批文档，然后使用向量搜索来对这批文档进行语义排序，提高精度
Java 代码实现 (线性加权 Rank Fusion):
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
public class HybridSearch {
private static final double BM25_WEIGHT = 0.6;
private static final double VECTOR_WEIGHT = 0.4;
public static void main(String[] args) {
List<String> documents = Arrays.asList(
"This is the first document about information retrieval.",
"This is the second document.",
"And this is the third one.",
"Is this the first document?"
);
String query = "first document retrieval";
// BM25
BM25 bm25 = new BM25(documents);
Map<Integer, Double> bm25Scores = new HashMap<>();
for (int i = 0; i < documents.size(); i++) {
bm25Scores.put(i, bm25.score(documents.get(i), query));
}
// Vector Search (使用简化版本)
List<double[]> documentVectors = Arrays.asList(
new double[]{0.1, 0.2, 0.3, 0.4},
new double[]{0.5, 0.6, 0.7, 0.8},
new double[]{0.9, 0.1, 0.2, 0.3},
new double[]{0.4, 0.5, 0.6, 0.7}
);
double[] queryVector = {0.2, 0.3, 0.4, 0.5};
Map<Integer, Double> vectorScores = new HashMap<>();
for (int i = 0; i < documentVectors.size(); i++) {
vectorScores.put(i, VectorSearch.cosineSimilarity(queryVector, documentVectors.get(i)));
}
// Rank Fusion (线性加权)
Map<Integer, Double> hybridScores = new HashMap<>();
for (int i = 0; i < documents.size(); i++) {
double bm25Score = bm25Scores.getOrDefault(i, 0.0);
double vectorScore = vectorScores.getOrDefault(i, 0.0);
hybridScores.put(i, BM25_WEIGHT * bm25Score + VECTOR_WEIGHT * vectorScore);
}
// 排序
List<Integer> rankedDocuments = hybridScores.entrySet().stream()
.sorted(Map.Entry.<Integer, Double>comparingByValue().reversed())
.map(Map.Entry::getKey)
.collect(Collectors.toList());
System.out.println("Ranked documents: " + rankedDocuments);
for (int docIndex : rankedDocuments) {
System.out.println("Document Index: " + docIndex + ", Document Content: " + documents.get(docIndex) + ", Hybrid Score: " + hybridScores.get(docIndex));
}
}
}这段代码演示了如何使用线性加权进行 Rank Fusion。 BM25_WEIGHT 和 VECTOR_WEIGHT 是可调参数，需要根据具体数据集进行调整。
Java 代码实现 (RRF Rank Fusion):
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
public class HybridSearchRRF {
private static final double K = 60; //调优参数
public static void main(String[] args) {
List<String> documents = Arrays.asList(
"This is the first document about information retrieval.",
"This is the second document.",
"And this is the third one.",
"Is this the first document?"
);
String query = "first document retrieval";
// BM25
BM25 bm25 = new BM25(documents);
Map<Integer, Double> bm25Scores = new HashMap<>();
for (int i = 0; i < documents.size(); i++) {
bm25Scores.put(i, bm25.score(documents.get(i), query));
}
List<Integer> bm25Ranked = bm25Scores.entrySet().stream()
.sorted(Map.Entry.<Integer, Double>comparingByValue().reversed())
.map(Map.Entry::getKey)
.collect(Collectors.toList());
// Vector Search (使用简化版本)
List<double[]> documentVectors = Arrays.asList(
new double[]{0.1, 0.2, 0.3, 0.4},
new double[]{0.5, 0.6, 0.7, 0.8},
new double[]{0.9, 0.1, 0.2, 0.3},
new double[]{0.4, 0.5, 0.6, 0.7}
);
double[] queryVector = {0.2, 0.3, 0.4, 0.5};
Map<Integer, Double> vectorScores = new HashMap<>();
for (int i = 0; i < documentVectors.size(); i++) {
vectorScores.put(i, VectorSearch.cosineSimilarity(queryVector, documentVectors.get(i)));
}
List<Integer> vectorRanked = vectorScores.entrySet().stream()
.sorted(Map.Entry.<Integer, Double>comparingByValue().reversed())
.map(Map.Entry::getKey)
.collect(Collectors.toList());
// Rank Fusion (Reciprocal Rank Fusion)
Map<Integer, Double> rrfScores = new HashMap<>();
for(int i = 0; i < documents.size(); i++){
rrfScores.put(i, 0.0);
}
for (int i = 0; i < bm25Ranked.size(); i++) {
int docId = bm25Ranked.get(i);
rrfScores.put(docId, rrfScores.get(docId) + 1.0 / (K + (i + 1)));
}
for (int i = 0; i < vectorRanked.size(); i++) {
int docId = vectorRanked.get(i);
rrfScores.put(docId, rrfScores.get(docId) + 1.0 / (K + (i + 1)));
}
// 排序
List<Integer> rankedDocuments = rrfScores.entrySet().stream()
.sorted(Map.Entry.<Integer, Double>comparingByValue().reversed())
.map(Map.Entry::getKey)
.collect(Collectors.toList());
System.out.println("Ranked documents: " + rankedDocuments);
for (int docIndex : rankedDocuments) {
System.out.println("Document Index: " + docIndex + ", Document Content: " + documents.get(docIndex) + ", RRF Score: " + rrfScores.get(docIndex));
}
}
}这段代码演示了如何使用 Reciprocal Rank Fusion进行Rank Fusion。 K是一个超参数，控制着排名对最终得分的影响。最佳的K值需要根据实际数据进行调整。
复杂业务场景匹配
在复杂的业务场景中，RAG 系统需要处理更复杂的查询和知识库。例如，查询可能涉及多个实体、关系和约束条件，知识库可能包含多种类型的数据和复杂的结构化信息。
为了应对这些挑战，我们可以采用以下策略：
- Query Decomposition (查询分解): 将复杂的查询分解为多个子查询，分别检索相关文档，然后将结果进行组合。
- Knowledge Graph Integration (知识图谱集成): 将知识图谱作为 RAG 系统的知识来源，利用知识图谱的结构化信息来增强检索和生成能力。
- Fine-tuning (微调): 使用特定领域的语料库对预训练模型进行微调，以提高其在特定领域的表现。
- Metadata Filtering (元数据过滤): 利用文档的元数据（例如，作者、日期、类别）来过滤检索结果，从而提高召回的准确性。
示例场景：保险索赔查询
假设我们有一个保险索赔的 RAG 系统，用户查询 "我的车在昨天下午三点被追尾了，对方全责，我应该怎么做？"。
- Query Decomposition: 将查询分解为 "车被追尾"、"对方全责"、"昨天下午三点" 等子查询。
- Knowledge Graph Integration: 利用知识图谱中的保险条款、索赔流程等信息来增强检索。
- Metadata Filtering: 根据事故发生的时间和地点，过滤相关的保险政策和理赔案例。
通过结合这些策略，我们可以构建更强大的 RAG 系统，能够更好地应对复杂业务场景的挑战。
评估指标
评估混合检索策略的性能至关重要。常用的评估指标包括：
- Precision (准确率): 召回的文档中相关文档的比例。
- Recall (召回率): 所有相关文档中被召回的文档的比例。
- F1-score: 准确率和召回率的调和平均值。
- Mean Average Precision (MAP): 多个查询的平均准确率的平均值。
- Normalized Discounted Cumulative Gain (NDCG): 考虑文档相关性等级的排序质量指标。
使用这些指标，我们可以比较不同混合检索策略的性能，并选择最佳策略。
需要注意的点
- 向量数据库的选择: 根据数据规模、查询性能和功能需求选择合适的向量数据库。Faiss、Annoy、Milvus 等都是不错的选择。
- 向量化模型的选择: 根据任务类型和数据特征选择合适的向量化模型。Sentence-BERT、BERT、GloVe 等都是常用的模型。
- 参数调优: 混合检索策略中的权重、参数等需要根据具体数据集进行调优，以获得最佳性能。
- 数据预处理: 对文本数据进行清洗、分词、去除停用词等预处理步骤，可以提高检索的准确性。
- 索引构建: 合理构建索引可以提高检索的效率。
总结
今天我们讨论了如何使用 Java 实现混合检索策略，特别是结合 BM25 和向量搜索，来提升 RAG 系统的召回精准度，并使其更好地适应复杂业务场景。我们深入探讨了 BM25 和向量搜索的原理及 Java 实现，以及如何使用 Rank Fusion 等方法将它们结合起来。此外，我们还讨论了如何应对复杂业务场景的挑战，并介绍了常用的评估指标。
通过灵活地运用这些技术，我们可以构建更强大的 RAG 系统，为用户提供更准确、更相关的答案，从而提升用户体验和业务价值。 记住，实际应用中，模型和参数都需要不断调整优化，以适应不同的数据集和业务需求。