5️⃣Mamba RAG with LangChain

Mamba를 이용한 검색 증강 세대(RAG) 구현과 Qdrant 데이터베이스 활용에 대해 집중적으로 살펴보고자 합니다. 최근 개발된 AI인 맘바는 특히 긴 시퀀스를 처리할 때 트랜스포머가 설정한 기존 표준에 도전합니다. RAG, Mamba, Qdrant의 시너지 효과는 효율성과 확장성의 강력한 조화를 약속하며, 대규모 데이터 처리 및 검색에 접근하는 방식을 혁신적으로 변화시킵니다.

맘바는 LSTM의 적응성과 상태 공간 모델의 효율성을 결합한 선택적 상태 공간이 돋보입니다. 전체 시퀀스를 한 번에 처리하는 기능은 트랜스포머를 연상시키지만 새로운 방식으로 변형되었습니다. 반면, RAG는 대규모 데이터 세트를 효율적으로 선별하고 정제하여 대규모 언어 모델(LLM)의 정밀도를 향상시키는 데 특화되어 있습니다.

Mamba in RAG

Mamba 아키텍처는 검색 증강 생성(RAG)의 기능을 강화하는 데 중추적인 역할을 합니다. 긴 시퀀스를 처리하는 혁신적인 접근 방식을 갖춘 맘바는 특히 RAG의 효율성과 정확성을 향상시키는 데 적합합니다. 선택적 상태 공간 모델은 기존의 상태 공간 모델에 비해 보다 유연하고 적응력 있는 상태 전환이 가능하기 때문에 RAG의 맥락에서 매우 효과적입니다.

  1. 긴 시퀀스 처리: 계산 효율성에 큰 타협 없이 긴 시퀀스로 확장할 수 있는 Mamba의 고유한 능력은 RAG에 매우 중요합니다. 이 특성은 광범위한 외부 지식 기반을 다룰 때 특히 유용하며, 검색 프로세스가 빠르고 정확하도록 보장합니다.

  2. 선택적 상태 공간: Mamba의 선택적 상태 공간은 시퀀스 처리에 대한 보다 미묘한 접근 방식을 제공합니다. 이 기능은 데이터베이스에서 검색된 쿼리와 해당 정보를 보다 동적이고 상황에 맞게 분석할 수 있게 해주므로 RAG의 컨텍스트 검색 프로세스에서 매우 유용합니다.

  3. 효율적인 계산: Mamba는 상태 공간 모델의 효율적인 계산 특성을 유지하여 전체 시퀀스의 포워드 패스를 한 번에 수행할 수 있습니다. 이러한 효율성은 특히 대량의 외부 데이터를 통합하고 처리할 때 RAG 프레임워크에서 유용합니다.

  4. 유연성과 적응성: Mamba의 아키텍처는 LSTM과 유사하게 시퀀스를 처리할 때 유연성과 적응성을 제공합니다. 이러한 유연성은 RAG에서 사용자 쿼리의 다양성과 예측 불가능성을 처리할 때 유리하며, 시스템이 광범위한 정보 검색 작업을 능숙하게 처리할 수 있도록 보장합니다.

기술 구현에 대해 자세히 알아보기 전에 검색 증강 세대(RAG), 맘바 아키텍처, Qdrant 데이터베이스의 개념을 실용적인 코드 중심 시나리오에서 어떻게 결합할 것인지에 대한 배경을 설정해 보겠습니다.

LangChain: RAG, Qdrant and Mamba

데이터 준비 및 처리, Mamba 모델 초기화 및 활용, 효율적인 데이터 검색 및 응답 생성을 위해 RAG와 Qdrant를 효과적으로 통합하는 방법을 구현해 보겠습니다. 다음 코드 분석은 프로세스의 각 단계에 대한 인사이트를 제공하며, Mamba의 계산 효율성과 Qdrant의 검색 기능의 시너지가 어떻게 RAG 기반 시스템의 성능을 향상시킬 수 있는지 보여줍니다.

1. Library Installation

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. Loading Data

RAG 프로세스에서 사용할 텍스트 문서가 포함된 데이터 세트(new_articles.zip)를 다운로드하고 압축을 풉니다.

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

3. Mamba Model Initialization

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. Tokenization

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 Process: Retrieval and Generation

이 스크립트에는 문서 로드, 텍스트 분할, Qdrant를 사용한 데이터베이스 색인 생성 함수가 포함되어 있습니다. RAG 프로세스의 중요한 단계인 관련 문서의 효율적인 벡터 기반 검색을 위해 Qdrant를 통합하는 방법을 보여줍니다.

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

그런 다음 predict 함수는 검색 부분과 생성 부분을 통합하여 검색된 문서가 제공하는 컨텍스트에 따라 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. Generating Responses

이 모델은 원래의 질문과 Qdrant 데이터베이스에서 검색된 컨텍스트를 모두 고려하여 사용자 쿼리("인생의 의미는 무엇인가요?", "판도는 얼마를 모금했나요?")에 대한 답변을 생성합니다. 이 단계는 Mamba의 효율적인 처리와 Qdrant의 검색 기능으로 강화된 RAG의 실제 적용을 보여줍니다.

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와 비교했을 때 속도뿐만 아니라 효율성 측면에서도 두드러지는데, 특히 크기가 더 작다는 점이 눈에 띕니다. 이러한 장점은 휴대폰이나 기타 저용량 디바이스와 같이 연산 능력이 제한적인 환경에 AI를 배포할 때 매우 유용할 수 있습니다.그러나 27억 개의 매개변수 모델은 더 큰 Transformer 모델에 비해 추론 능력이 약간 뒤처지는 것으로 보입니다. 앞으로 추론 능력을 향상시키기 위해 모델을 미세 조정하는 것은 가치 있는 단계가 될 수 있습니다. 하지만 현재로서는 성능과 효율성이 균형을 이루고 있어 특히 컴퓨팅 리소스가 제약이 있는 애플리케이션에서 매력적인 선택이 될 수 있습니다. 이 모델은 고급 AI 기술의 접근성과 적용 가능성을 넓힐 수 있는 잠재력을 지니고 있습니다.

PreviousMambaNextLlama3-8B with LlamaIndex

Last updated