Mistral과 Gradio로 PDF와 채팅하기
Mistral과 Gradio로 PDF와 채팅하기 (Chat with Your PDF using Mistral and Gradio)
gradio를 사용해 채팅과 PDF 읽기 기능을 갖춘 챗봇을 구축하는 기본 방법을 배우는 문서예요. 진보된 RAG(검색 증강 생성) 기법을 활용해 PDF 내용에 대해 모델과 대화할 수 있습니다.
출처: 문서
본문
이 가이드에서는 gradio로 채팅과 PDF 읽기 기능을 갖춘 챗봇을 구축하는 기본을 소개할게요.
채팅 인터페이스 (Chat Interface)
먼저 간단한 채팅 인터페이스를 구현해요. 이를 위해 gradio, mistralai 라이브러리와 mistralai.models.chat_completion의 ChatMessage를 임포트해야 해요.
pip install gradio mistralai
이 데모는 gradio===4.32.2와 mistralai===0.4.0을 사용해요.
import gradio as gr
from mistralai.client import MistralClient
from mistralai.models.chat_completion import ChatMessage
다음으로 Mistral API 키로 MistralClient 인스턴스를 만들어요.
mistral_api_key = "your_api_key"
cli = MistralClient(api_key = mistral_api_key)
gradio의 ChatInterface를 사용해 인터페이스를 만들 수 있어요. 다음과 같이 생겼어요:
def ask_mistral(message: str, history: list):
return "Bot's response."
app = gr.ChatInterface(fn = ask_mistral, title = "Ask Mistral")
app.launch()
이제 ask_mistral을 수정해 메시지와 히스토리를 파싱하고 Mistral API를 호출하며 응답을 스트리밍하기만 하면 돼요.
def ask_mistral(message: str, history: list):
messages = []
for couple in history:
messages.append(ChatMessage(role = "user", content = couple[0]))
messages.append(ChatMessage(role = "assistant", content = couple[1]))
messages.append(ChatMessage(role = "user", content = message))
full_response = ""
for chunk in cli.chat_stream(model = "open-mistral-7b", messages = messages, max_tokens = 1024):
full_response += chunk.choices[0].delta.content
yield full_response
완료! 준비되면 스크립트(chat.py)를 실행할 수 있어요.
import gradio as gr
from mistralai.client import MistralClient
from mistralai.models.chat_completion import ChatMessage
mistral_api_key = "your_api_key"
cli = MistralClient(api_key = mistral_api_key)
def ask_mistral(message: str, history: list):
messages = []
for couple in history:
messages.append(ChatMessage(role = "user", content = couple[0]))
messages.append(ChatMessage(role = "assistant", content = couple[1]))
messages.append(ChatMessage(role = "user", content = message))
full_response = ""
for chunk in cli.chat_stream(model = "open-mistral-7b", messages = messages, max_tokens = 1024):
full_response += chunk.choices[0].delta.content
yield full_response
app = gr.ChatInterface(fn = ask_mistral, title = "Ask Mistral")
app.launch()
PDF와 채팅하기 (Chatting with PDFs)
모델이 PDF를 읽게 하려면 콘텐츠를 변환하고 텍스트를 추출한 다음, Mistral의 임베딩 모델로 문서(들)의 청크를 검색해 모델에 공급해야 해요. 기본적인 RAG를 구현해야 하죠!
이 작업에는 faiss, PyPDF2와 다른 라이브러리가 필요해요. 임포트해 볼게요.
pip install numpy PyPDF2 faiss
CPU만 사용한다면 faiss-cpu를 대신 설치하세요. 이 데모는 numpy===1.26.4, PyPDF2===0.4.0, faiss-cpu===1.8.0을 사용해요.
import gradio as gr
from mistralai.client import MistralClient
from mistralai.models.chat_completion import ChatMessage
import numpy as np
import PyPDF2
import faiss
인터페이스가 파일 업로드를 허용하게 하려면 ChatInterface에서 멀티모달(multimodal)을 켜야 해요.
app = gr.ChatInterface(fn = ask_mistral, title = "Ask Mistral and talk to your PDFs", multimodal = True)
app.launch()
이제 인터페이스가 파일도 받아줘요. 다음 단계는 파일을 처리하고 메시지에서 PDF 파일을 걸러내는 거예요.
def ask_mistral(message: str, history: list):
messages = []
pdfs = message["files"]
for couple in history:
if type(couple[0]) is tuple:
pdfs += couple[0]
else:
messages.append(ChatMessage(role = "user", content = couple[0]))
messages.append(ChatMessage(role = "assistant", content = couple[1]))
messages.append(ChatMessage(role = "user", content = message["text"]))
full_response = ""
for chunk in cli.chat_stream(model = "open-mistral-7b", messages = messages, max_tokens = 1024):
full_response += chunk.choices[0].delta.content
yield full_response
PDF 파일을 읽고 RAG를 구현할 준비가 됐어요. 이를 위해 PDF에서 관련 텍스트 청크를 단일 문자열로 연결해 검색하는 함수를 만들어야 해요. Mistral의 임베딩을 사용할 거예요. 텍스트를 임베딩으로 변환하는 함수를 빠르게 설계해 볼게요.
def get_text_embedding(input: str):
embeddings_batch_response = cli.embeddings(
model = "mistral-embed",
input = input
)
return embeddings_batch_response.data[0].embedding
이제 모든 RAG를 처리하고 적절한 청크를 검색하는 rag_pdf를 만들 수 있어요.
def rag_pdf(pdfs: list, question: str) -> str:
chunk_size = 4096
chunks = []
for pdf in pdfs:
chunks += [pdf[i:i + chunk_size] for i in range(0, len(pdf), chunk_size)]
text_embeddings = np.array([get_text_embedding(chunk) for chunk in chunks])
d = text_embeddings.shape[1]
index = faiss.IndexFlatL2(d)
index.add(text_embeddings)
question_embeddings = np.array([get_text_embedding(question)])
D, I = index.search(question_embeddings, k = 4)
retrieved_chunk = [chunks[i] for i in I.tolist()[0]]
text_retrieved = "\n\n".join(retrieved_chunk)
return text_retrieved
이 함수에서는 PDF 파일을 동일한 크기의 청크로 자르고, 그 임베딩을 얻어 faiss로 벡터 검색을 적용해 최고의 4개 청크를 검색해요. 다음이자 마지막 단계는 PyPDF2로 PDF 파일 자체를 읽고 모델과 통합하는 거예요.
def ask_mistral(message: str, history: list):
messages = []
pdfs = message["files"]
for couple in history:
if type(couple[0]) is tuple:
pdfs += couple[0]
else:
messages.append(ChatMessage(role = "user", content = couple[0]))
messages.append(ChatMessage(role = "assistant", content = couple[1]))
if pdfs:
pdfs_extracted = []
for pdf in pdfs:
reader = PyPDF2.PdfReader(pdf)
txt = ""
for page in reader.pages:
txt += page.extract_text()
pdfs_extracted.append(txt)
retrieved_text = rag_pdf(pdfs_extracted, message["text"])
messages.append(ChatMessage(role = "user", content = retrieved_text + "\n\n" + message["text"]))
else:
messages.append(ChatMessage(role = "user", content = message["text"]))
full_response = ""
for chunk in cli.chat_stream(model = "open-mistral-7b", messages = messages, max_tokens = 1024):
full_response += chunk.choices[0].delta.content
yield full_response
이제 준비 완료! 스크립트 chat_with_pdfs.py를 실행할 수 있어요.
import gradio as gr
from mistralai.client import MistralClient
from mistralai.models.chat_completion import ChatMessage
import numpy as np
import PyPDF2
import faiss
mistral_api_key = "your_api_key"
cli = MistralClient(api_key = mistral_api_key)
def get_text_embedding(input: str):
embeddings_batch_response = cli.embeddings(
model = "mistral-embed",
input = input
)
return embeddings_batch_response.data[0].embedding
def rag_pdf(pdfs: list, question: str) -> str:
chunk_size = 4096
chunks = []
for pdf in pdfs:
chunks += [pdf[i:i + chunk_size] for i in range(0, len(pdf), chunk_size)]
text_embeddings = np.array([get_text_embedding(chunk) for chunk in chunks])
d = text_embeddings.shape[1]
index = faiss.IndexFlatL2(d)
index.add(text_embeddings)
question_embeddings = np.array([get_text_embedding(question)])
D, I = index.search(question_embeddings, k = 4)
retrieved_chunk = [chunks[i] for i in I.tolist()[0]]
text_retrieved = "\n\n".join(retrieved_chunk)
return text_retrieved
def ask_mistral(message: str, history: list):
messages = []
pdfs = message["files"]
for couple in history:
if type(couple[0]) is tuple:
pdfs += couple[0]
else:
messages.append(ChatMessage(role= "user", content = couple[0]))
messages.append(ChatMessage(role= "assistant", content = couple[1]))
if pdfs:
pdfs_extracted = []
for pdf in pdfs:
reader = PyPDF2.PdfReader(pdf)
txt = ""
for page in reader.pages:
txt += page.extract_text()
pdfs_extracted.append(txt)
retrieved_text = rag_pdf(pdfs_extracted, message["text"])
messages.append(ChatMessage(role = "user", content = retrieved_text + "\n\n" + message["text"]))
else:
messages.append(ChatMessage(role = "user", content = message["text"]))
full_response = ""
for chunk in cli.chat_stream(model = "open-mistral-7b", messages = messages, max_tokens = 1024):
full_response += chunk.choices[0].delta.content
yield full_response
app = gr.ChatInterface(fn = ask_mistral, title = "Ask Mistral and talk to your PDFs", multimodal = True)
app.launch()