使用Mamba和Qdrant数据库实现RAG:详细探索

2024年02月04日 由 alex 发表 493 0

1


大家好,今天,我们将深入探讨人工智能的迷人世界,尤其是通过 Mamba 实现检索增强生成(RAG)和利用 Qdrant 数据库。Mamba 是人工智能领域的最新发展,它挑战了变形金刚所设定的传统规范,尤其是在处理冗长序列方面。RAG、Mamba 和 Qdrant 的协同作用有望实现效率和可扩展性的完美融合,彻底改变我们处理大规模数据和检索的方式。


了解RAG和Mamba 的基础知识: 概述


Mamba 凭借其选择性状态空间脱颖而出,融合了 LSTM 的适应性和状态空间模型的效率。其一次性处理整个序列的能力让人联想起变形金刚,但又不失新意。另一方面,RAG 擅长通过高效筛选和完善海量数据集来提高大型语言模型(LLM)的精确度。


Mamba在RAG中的作用


Mamba 架构在增强检索增强生成(RAG)功能方面发挥着关键作用。Mamba 采用创新方法处理冗长序列,特别适合提高 RAG 的效率和准确性。与传统的状态空间模型相比,它的选择性状态空间模型允许更灵活、适应性更强的状态转换,因此在 RAG 中非常有效。


Mamba如何改进RAG


  1. 处理长序列: Mamba 能够在不明显降低计算效率的情况下扩展到更长的序列,这一固有能力对 RAG 至关重要。在处理大量外部知识库时,这一特性尤为有利,可确保检索过程既快速又准确。
  2. 选择性状态空间:Mamba 中的选择性状态空间为序列处理提供了一种更加细致入微的方法。这一功能在 RAG 的上下文检索过程中非常宝贵,因为它允许对查询和从数据库中检索的相应信息进行更动态和对上下文更敏感的分析。
  3. 高效计算: Mamba 保留了状态空间模型的高效计算特性,使其能够在一次扫描中对整个序列进行前向传递。这种效率在 RAG 框架中非常有用,尤其是在整合和处理大量外部数据时。
  4. 灵活性和适应性: Mamba 的架构类似于 LSTM,在处理序列时具有灵活性和适应性。这种灵活性在处理 RAG 中用户查询的多样性和不可预测性时非常有利,可确保系统能够熟练处理各种信息检索任务。


在深入探讨技术实现之前,让我们先来了解一下我们是如何将检索增强生成(RAG)、Mamba 架构和 Qdrant 数据库的概念结合到实际的代码驱动方案中的。


如何利用Mamba模型并将RAG与Qdrant整合以实现高效数据检索


在本节中,我们将探讨一个 Python 脚本,它是这些先进技术集成的范例。该脚本不仅说明了必要环境和库的安装和设置,还演示了如何准备和处理数据、初始化和使用 Mamba 模型,以及如何将 RAG 与 Qdrant 有效集成,以实现高效的数据检索和响应生成。下面的代码分解将深入介绍流程的每个步骤,展示 Mamba 的计算效率和 Qdrant 的检索能力如何协同提高基于 RAG 的系统性能。


1. 环境设置和库安装


首先,脚本会安装必要的库,包括 PyTorch、Mamba-SSM、LangChain、Qdrant 客户端等。这些安装对于设置 RAG、Mamba 和 Qdrant 协同工作所需的环境至关重要。


from inspect import cleandoc
import pandas as pd
import torch
from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
from transformers import AutoTokenizer
from langchain.embeddings import HuggingFaceBgeEmbeddings
from langchain_community.vectorstores import Qdrant
from langchain.document_loaders import TextLoader
from langchain.document_loaders import DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter


2. 加载数据


然后,它下载并解压一个数据集(new_articles.zip),该数据集可能包含将在 RAG 过程中使用的文本文件。


!wget -q https://www.dropbox.com/s/vs6ocyvpzzncvwh/new_articles.zip
!unzip -q new_articles.zip -d new_articles


3.Manba模型初始化


Mamba 模型以特定的模型名称("havenhq/mamba-chat")进行初始化,并根据可用性设置为使用 GPU 或 CPU。这一步对于在 RAG 中利用 Mamba 高效计算长序列至关重要。


MODEL_NAME = "havenhq/mamba-chat"
model = MambaLMHeadModel.from_pretrained(MODEL_NAME, device=DEVICE, dtype=torch.float16)


4. 标记化和模型输入准备


标记化器为模型准备输入。它被配置为适当处理 Mamba 模型的输入和输出。标记化过程对于将用户查询转换为 Mamba 可以处理的格式至关重要。


ANSWER_START = "<|assistant|>\n"
ANSWER_END = "<|endoftext|>"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
tokenizer.eos_token = ANSWER_END
tokenizer.pad_token = tokenizer.eos_token
tokenizer.chat_template = AutoTokenizer.from_pretrained(
    "BAAI/bge-small-en-v1.5"
).chat_template


5. RAG 流程: 检索和生成


脚本包括使用 Qdrant 加载文档、分割文本和创建数据库索引的功能。它说明了 Qdrant 与相关文档的高效矢量检索的集成,这是 RAG 流程的关键步骤。


loader = DirectoryLoader('./new_articles/', glob="./*.txt", loader_cls=TextLoader)
documents = loader.load()
#splitting the text into
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
texts = text_splitter.split_documents(documents)
def get_index(): #creates and returns an in-memory vector store to be used in the application
    model_name = "BAAI/bge-small-en-v1.5"
    encode_kwargs = {'normalize_embeddings': True} # set True to compute cosine similarity
    embeddings = HuggingFaceBgeEmbeddings(
    model_name=model_name,
    model_kwargs={'device': 'cpu'},
    encode_kwargs=encode_kwargs
        )
    index_from_loader = Qdrant.from_documents(
            texts,
            embeddings,
            location=":memory:",  # Local mode with in-memory storage only
            collection_name="my_documents",
        )
    return index_from_loader #return the index to be cached by the client app
vector_index = get_index()


Qdrant 可作为我们的矢量数据库,因为它具有快速索引和查询功能,并支持各种距离度量。这使它成为管理大量矢量数据的理想选择,并提高了搜索的准确性和相关性。


语义搜索(semantic_search)函数执行 RAG 的检索部分,通过查询 Qdrant 向量索引来查找与给定提示相关的文档。


def semantic_search(index, original_prompt): #rag client function
    relevant_prompts = index.similarity_search(original_prompt)
    list_prompts = []
    for i in range(len(relevant_prompts)):
        list_prompts.append(relevant_prompts[i].page_content)
    return list_prompts


然后,预测功能将检索部分与生成部分整合在一起,在生成部分,Mamba 模型根据检索文档提供的上下文生成响应。


def predict(prompt: str) -> str:
    selected_prompt = semantic_search(vector_index, prompt)
    selected_prompt = ' , '.join(selected_prompt)
    messages = []
    if selected_prompt:
        messages.append({"role": "system", "content": "Please respond to the original query. If the selected document prompt is relevant and informative, provide a detailed answer based on its content. However, if the selected prompt does not offer useful information or is not applicable, simply state 'No answer found'."})
    messages.append({"role": "user", "content": f"""Original Prompt: {prompt}\n\n
                    Selected Prompt: {selected_prompt}\n\n
                    respond: """})
    input_ids = tokenizer.apply_chat_template(
        messages, return_tensors="pt", add_generation_prompt=True
    ).to(DEVICE)
    outputs = model.generate(
        input_ids=input_ids,
        max_length=1024,
        temperature=0.9,
        top_p=0.7,
        eos_token_id=tokenizer.eos_token_id,
    )
    response = tokenizer.decode(outputs[0])
    return extract_response(response)


6. 生成回复


该模型通过考虑原始提示和从 Qdrant 数据库检索的上下文来生成对用户查询 ( "What is the meaning of life?"、 ) 的响应。"How much money did Pando raise?"这一步展示了 RAG 的实际应用,并通过 Mamba 的高效处理和 Qdrant 的检索能力得到增强。


predict("How much money did Pando raise?")
>>> """
Selected Prompt: How much money did Pando raise?\n\nSelected Answer: $30 million in a Series B round, bringing its total raised to $45 million.
"""
predict("What is the news about Pando?")
>>>"""    
Selected Prompt: What is the news about Pando?\n\nSelected Response: Pando has raised $30 million in a Series B round, bringing its total raised to $45 million. The startup is led by Nitin Jayakrishnan and Abhijeet Manohar, who previously worked together at iDelivery, an India-based freight tech marketplace. The startup is focused on global logistics and supply chain management through a software-as-a-service platform. Pando has a compelling sales, marketing and delivery capabilities, according to Jayakrishnan. The startup has also tapped existing enterprise users at warehouses, factories, freight yards and ports and expects to expand its customer base. The company is also open to exploring strategic partnerships and acquisitions with this round of funding.
"""


在当前的实验中,我们使用了一个拥有 27 亿个参数的模型,观察它的性能非常有趣。值得注意的是,它的运行效率几乎与 70 亿参数的 LLaMA2 模型相当。与 LLaMA2-7B 模型相比,它不仅在速度上,而且在效率上都更胜一筹。在计算能力有限的环境(如手机或其他低容量设备)中部署人工智能时,这一优势至关重要。


不过,有得必有失;与一些较大的变形金刚相比,27 亿参数模型的推理能力似乎略显落后。展望未来,对模型进行微调以增强其推理能力可能是一个有价值的步骤。不过,就目前而言,它在性能和效率之间的平衡使其成为一个令人信服的选择,尤其是在计算资源有限的应用领域。该模型有望扩大先进人工智能技术的可及性和适用性。


总结


总之,RAG、Mamba 和 Qdrant 的集成证明了人工智能领域对创新的不懈追求。它标志着人工智能朝着更高效、更易用、更有能力处理数字世界中日益增长的数据处理需求的方向迈出了一步。

文章来源:https://medium.com/ai-insights-cobet/implementing-rag-with-mamba-and-the-qdrant-database-a-detailed-exploration-with-code-3e9a12b610f3
欢迎关注ATYUN官方公众号
商务合作及内容投稿请联系邮箱:bd@atyun.com
评论 登录
写评论取消
回复取消