Mesop으로 PDF와 채팅하기

Mesop으로 PDF와 채팅하기 (Mistral)

mesop를 사용해 채팅과 PDF 읽기 기능을 갖춘 챗봇을 만드는 기본기를 소개하는 가이드예요. Mistral로 대화 인터페이스를 만들고, PDF에서 텍스트를 추출해 간단한 RAG로 문서와 대화할 수 있게 합니다.

출처: 문서

본문

채팅 인터페이스 (Chat Interface)

먼저 간단한 채팅 인터페이스를 구현합니다. 이를 위해 mesop, mesop.labs, mistralai 라이브러리와 mistralai.models.chat_completion의 ChatMessage가 필요해요.

pip install mesop mistralai

이 데모는 mesop===0.9.3과 mistralai===0.4.0을 사용합니다.

import mesop as me
import mesop.labs as mel
from mistralai.client import MistralClient
from mistralai.models.chat_completion import ChatMessage

Mistral API 키로 MistralClient 인스턴스를 만듭니다.

mistral_api_key = "api_key"
cli = MistralClient(api_key = mistral_api_key)

mesop로 인터페이스를 만들 때는 mel.chat 함수를 이용합니다.

def ask_mistral(message: str, history: list[mel.ChatMessage]):
    messages = [ChatMessage(role=m.role, content=m.content) for m in history[:-1]]
    for chunk in cli.chat_stream(model = "open-mistral-7b", messages = messages, max_tokens = 1024):
        yield chunk.choices[0].delta.content

@me.page(title="Talk to Mistral")
def page():
    mel.chat(ask_mistral, title="Ask Mistral", bot_user="Mistral")

이제 mesop chat.py 명령만 실행하면 됩니다!

PDF와 채팅하기 (Chatting with PDFs)

모델이 PDF를 읽게 하려면 내용을 변환·텍스트 추출하고, Mistral의 임베딩 모델을 사용해 문서의 청크를 검색해 모델에 제공해야 해요. 간단한 RAG(Retrieval-Augmented Generation)를 구현하는 셈이죠.

이 작업에는 faiss, PyPDF2와 그 외 라이브러리가 필요합니다.

pip install numpy PyPDF2 faiss

CPU 전용이라면 faiss-cpu를 설치하세요. 이 데모는 numpy===1.26.4, PyPDF2===0.4.0, faiss-cpu===1.8.0을 사용합니다.

import io
import mesop as me
import mesop.labs as mel
from mistralai.client import MistralClient
from mistralai.models.chat_completion import ChatMessage
import numpy as np
import PyPDF2
import faiss

인터페이스에서 파일 업로드를 지원하려면 page 함수에 업로더를 추가해야 해요.

@me.page(title="Talk to Mistral")
def page():
    with me.box(style=me.Style(height = "100%", display="flex", flex_direction="column", align_items="center",padding=me.Padding(top = 0, left = 30, right = 30, bottom = 0))):
        with me.box(style=me.Style(padding=me.Padding(top = 16), position="fixed")):
            me.uploader(
                label="Upload PDF",
                accepted_file_types=["file/pdf"],
                on_upload=handle_upload,
            )
        with me.box(style=me.Style(width="100%")):
            mel.chat(ask_mistral, title="Ask Mistral", bot_user="Mistral")

이제 인터페이스가 파일도 받습니다. 다음 단계는 파일을 처리하고 PDF에서 텍스트를 추출하는 것입니다.

@me.stateclass
class State:
    content: str

def handle_upload(event: me.UploadEvent):
    state = me.state(State)
    reader = PyPDF2.PdfReader(io.BytesIO(event.file.getvalue()))
    txt = ""
    for page in reader.pages:
        txt += page.extract_text()
    state.content = txt

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개 청크를 가져옵니다. 마지막 단계는 이를 모델과 통합하는 것이에요.

def ask_mistral(message: str, history: list[mel.ChatMessage]):
    messages = [ChatMessage(role=m.role, content=m.content) for m in history[:-1]]

    state = me.state(State)
    if state.content:
        retrieved_text = rag_pdf([state.content], message)
        messages[-1] = ChatMessage(role = "user", content = retrieved_text + "\n\n" +messages[-1].content)

    for chunk in cli.chat_stream(model = "open-mistral-7b", messages = messages, max_tokens = 1024):
        yield chunk.choices[0].delta.content

이제 준비가 끝났어요! mesop chat_with_pdfs.py 명령으로 스크립트를 실행할 수 있습니다.

더 알아보기 (Learn more)

Mesop 공식 문서 — Google의 Python 웹 UI 프레임워크
mel.chat / me.uploader — Mesop의 채팅·파일 업로드 컴포넌트
faiss.IndexFlatL2 — 유사도 검색용 FAISS 인덱스
모델: open-mistral-7b, mistral-embed