LangChain混合检索器优化：降低延迟与Token消耗

2025-02-27 03:59:50

LangChain 混合检索器性能优化：降低延迟与 Token 消耗

使用 LangChain 构建检索增强生成 (RAG) 应用时，混合检索器是个常用的东西。这哥们结合了关键词匹配和语义相似性，检索效果通常不错。但是！性能问题很烦人，经常遇到延迟高、返回的 token 太多这种破事儿，既影响速度，又烧钱。

咱们先来看看问题出在哪儿。

问题根源分析

你当前的设置是这样的：

from langchain.retrievers import EnsembleRetriever, create_tfidf_retriever
from langchain.vectorstores import VectorStore

# 初始化 TF-IDF 检索器，k=5
tfidf_retriever = create_tfidf_retriever(documents=documents)
tfidf_retriever.k = 5

# 初始化密集检索器，使用向量存储
dense_retriever = vectorstore.as_retriever(search_kwargs={"k": 5})

# 组合检索器，权重均衡
ensemble_retriever = EnsembleRetriever(
    retrievers=[dense_retriever, tfidf_retriever],
    weights=[0.5, 0.5]
)

这种设置能检索到相关结果，但速度慢，而且经常返回一堆用不上的文本，token 数超标。就算你尝试用 MultiQueryRetriever 来处理改述问题，虽然相关性提高了，但延迟也更高了。

归根结底，几个原因：

检索数量过多（k 值）： 两个检索器的 k 值都设置为 5，加起来可能返回最多 10 个文档（虽然有去重，但数量还是不少）。
权重分配： 50/50 的权重不一定是最优的。有时候关键词匹配更重要，有时候语义相似性更重要。
MultiQueryRetriever 的开销： 这个检索器会生成多个查询，虽然提高了多样性，但也增加了 API 调用次数和计算量。
** 缺少后处理过滤步骤** 返回了太多和问题无关紧要的内容, 可以过滤一次.

优化方案：各个击破

针对这些问题，咱们一个一个解决。

1. 调整 k 值和权重

最直接的方法就是调整每个检索器的 k 值和组合检索器的权重。

原理： k 值决定了每个检索器返回的文档数量。权重决定了最终结果中，每个检索器的结果占比。

代码示例：

# 降低 k 值，减少返回文档数量
tfidf_retriever.k = 3
dense_retriever = vectorstore.as_retriever(search_kwargs={"k": 3})

# 调整权重，偏向 TF-IDF 或密集检索
ensemble_retriever = EnsembleRetriever(
    retrievers=[dense_retriever, tfidf_retriever],
    weights=[0.7, 0.3]  # 偏向密集检索
)

或者

ensemble_retriever = EnsembleRetriever(
  retrievers=[dense_retriever, tfidf_retriever],
  weights=[0.3, 0.7]  # 偏向TF-IDF检索
)

操作建议：
- 根据你的数据集和查询特点，多试几组 k 值和权重。
- 如果查询通常比较具体，可以适当增加 TF-IDF 的权重。
- 如果查询通常比较抽象，可以适当增加密集检索的权重。
进阶技巧 : 可以写一个测试脚本, 测试集,通过程序自动化尝试不同的K和权重值,找到最符合需求的配置

2. 使用更高效的向量数据库和 Embedding 模型

向量数据库和 Embedding 模型的选择，直接影响密集检索的速度。

原理： 向量数据库负责存储和检索向量。Embedding 模型负责将文本转换为向量。更高效的数据库和模型，检索速度更快。
推荐：
- 向量数据库： 考虑使用 FAISS、Annoy 或 ScaNN 等专门为向量检索优化的数据库。这些数据库通常比通用数据库（如 PostgreSQL）快得多。如果文档有ID,优先使用FAISS 的IndexIDMap索引.
- Embedding 模型： 考虑使用 Sentence Transformers 或 Instructor Embeddings 等模型。这些比简单的 word embeddings（如 Word2Vec 或 GloVe）生成的向量更能表达句子的语义。
** 建议: ** 优先使用比较新的 Embedding 模型，效果一般比老的模型效果要好很多。

代码示例 (使用 FAISS)：

from langchain.vectorstores import FAISS
from langchain.embeddings import HuggingFaceInstructEmbeddings #或 HuggingFaceEmbeddings

 # 使用 Instructor Embeddings 或 Sentence Transformers 模型
embeddings = HuggingFaceInstructEmbeddings(model_name="hkunlp/instructor-large")
#或者使用 sentence-transformer 的 all-mpnet-base-v2
#embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")

# 从文档创建 FAISS 索引, 并指定ID
db = FAISS.from_documents(docs, embeddings,index_factory_string="IndexFlatL2")#IDMap,Flat")

dense_retriever = db.as_retriever(search_kwargs={"k": 3})

3. 结果后处理：精简 Token

就算检索到的文档相关，也可能包含很多冗余信息。可以在将文档传递给 LLM 之前，对它们进行处理。

原理： 通过删除不相关的句子、段落或关键词，减少 token 数量。
方法：
- 基于相似度： 计算文档中每个句子与查询的相似度，只保留相似度高的句子。
- 基于关键词： 提取文档中的关键词，只保留包含关键词的句子。
- MMR（最大边际相关性）： 这是一种更高级的方法，它在选择句子时，不仅考虑与查询的相关性，还考虑句子之间的多样性，避免选择重复的句子。

代码示例 (基于相似度)：

from sklearn.metrics.pairwise import cosine_similarity
from sentence_transformers import SentenceTransformer

# 使用 Sentence Transformers 模型计算相似度
model = SentenceTransformer('all-MiniLM-L6-v2')

def filter_sentences(query, sentences, threshold=0.6):
    query_embedding = model.encode(query)
    sentence_embeddings = model.encode(sentences)
    similarities = cosine_similarity([query_embedding], sentence_embeddings)[0]
    filtered_sentences = [sentences[i] for i, score in enumerate(similarities) if score >= threshold]
    return filtered_sentences

#检索到的文档, 可能是多个.  
documents = ensemble_retriever.get_relevant_documents("...你的查询...")

filtered_docs = []

for doc in documents:    
    sentences = doc.page_content.split('。')  # 简单按句号分割
    filtered = filter_sentences("...你的查询...", sentences)
    doc.page_content = '。'.join(filtered)  # 用句号重新连接
    filtered_docs.append(doc)

注意 : 这个过程也要消耗一定的时间，具体实现需要考虑平衡效果和性能消耗。

4. 调整 `MultiQueryRetriever`

MultiQueryRetriever 提高了检索多样性，但也增加了延迟。可以对它进行微调。

原理： MultiQueryRetriever 使用 LLM 生成多个查询变体。可以调整生成查询的数量和使用的 LLM。

代码示例：

from langchain.retrievers.multi_query import MultiQueryRetriever
from langchain.chat_models import ChatOpenAI

# 使用更快的 LLM (如果可以的话)
llm = ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo") #例如使用gpt-3.5-turbo

retriever_from_llm = MultiQueryRetriever.from_llm(
    retriever=ensemble_retriever,
    llm=llm,
    parser_key="lines" #或 "text", 根据你的 LLM 输出格式选择
)

 #减少生成查询的数量，这样会变快
retriever_from_llm.query_count = 3

或者干脆避免使用 MultiQueryRetriever. 直接提供明确的提示，帮助用户提问。

5. 缓存

对于相同的查询，可以缓存检索结果，避免重复计算。

原理： 将查询和对应的检索结果存储起来。下次遇到相同查询时，直接返回缓存的结果。

代码示例 (使用 Python 内置的 lru_cache)：

from functools import lru_cache

@lru_cache(maxsize=128)  # 缓存最多 128 个查询
def get_relevant_documents_cached(query):
    return ensemble_retriever.get_relevant_documents(query)

#第一次执行会检索，并把检索结果存在cache里。
docs=get_relevant_documents_cached("...你的查询...")    
# 第二次相同的查询会立即从缓存返回。
docs=get_relevant_documents_cached("...你的查询...")