Mistral AI와 Pinecone으로 RAG 구축하기

Mistral의 임베딩 모델 mistral-embed와 벡터 데이터베이스 Pinecone을 사용해 RAG 파이프라인을 만드는 쿡북이에요. ArXiv 논문 데이터를 임베딩해 Pinecone에 저장하고, 검색한 뒤 Mistral Large로 답변을 생성합니다.

출처: 문서

본문

데이터 준비 (Data Preparation)

먼저 필요한 라이브러리를 설치합니다.

!pip3 install -qU datasets mistralai pinecone

인코딩해서 저장할 데이터셋을 내려받습니다. jamescalam/ai-arxiv2-semantic-chunks 데이터셋은 LLM과 GenAI 중심의 인기 ArXiv 논문에서 스크랩한 데이터를 담고 있어요.

from datasets import load_dataset

data = load_dataset(
    "jamescalam/ai-arxiv2-semantic-chunks",
    split="train[:10000]"
)
data

약 20만 개의 청크가 있으며, 각 청크는 대략 1~2문단 길이입니다.

data[0]

필요한 형식으로 데이터를 정리합니다. id, 임베딩할 text, 그리고 메타데이터를 담습니다.

data = data.map(lambda x: {
    "id": x["id"],
    "metadata": {
        "title": x["title"],
        "content": x["content"],
    }
})

# drop unneeded columns
data = data.remove_columns([
    "title", "content", "prechunk_id",
    "postchunk_id", "arxiv_id", "references"
])
data

검색을 위한 벡터를 만들 임베딩 모델로 Mistral AI의 mistral-embed를 사용합니다. 이 모델은 비용이 들 수 있으니 유의하세요 (이 노트북 실행 비용은 <$1).

import os
from mistralai.client import Mistral
import getpass  # console.mistral.ai/api-keys/
# get API key from left navbar in Mistral console
mistral_api_key = os.getenv("MISTRAL_API_KEY") or getpass.getpass("Enter your Mistral API key: ")

# initialize client
mistral = Mistral(api_key=mistral_api_key)

임베딩을 생성합니다.

embed_model = "mistral-embed"
embeds = mistral.embeddings.create(
    model=embed_model, inputs=["this is a test"]
)

반환된 임베딩의 차원을 확인합니다. 벡터 인덱스를 초기화할 때 이 값이 필요해요.

dims = len(embeds.data[0].embedding)
dims

벡터 DB인 Pinecone을 생성합니다. 이를 위해 무료 Pinecone API key가 필요해요. API key는 Pinecone 대시보드 왼쪽 내비게이션의 "API Keys" 버튼에서 찾을 수 있습니다.

from pinecone import Pinecone
# initialize connection to pinecone (get API key at app.pinecone.io)
api_key = os.getenv("PINECONE_API_KEY") or getpass.getpass("Enter your Pinecone API key: ")

# configure client
pc = Pinecone(api_key=api_key)

인덱스 사양을 설정해 클라우드 제공자와 지역을 정합니다. 사용 가능한 모든 제공자·지역 목록은 여기에서 확인할 수 있어요.

from pinecone import ServerlessSpec

spec = ServerlessSpec(
    cloud="aws", region="us-east-1"
)

인덱스를 생성합니다. mistral-embed의 차원(1024)과 호환되는 metric(cosine 또는 dotproduct)을 사용합니다.

import time

index_name = "mistral-rag"

existing_indexes = [
    index_info["name"] for index_info in pc.list_indexes()
]

# check if index already exists (it shouldn't if this is first time)
if index_name not in existing_indexes:
    # if does not exist, create index
    pc.create_index(
        index_name,
        dimension=dims,  # dimensionality of mistral-embed
        metric='dotproduct',
        spec=spec
    )

# wait for index to be initialized
while not pc.describe_index(index_name).status['ready']:
    time.sleep(1)

# connect to index
index = pc.Index(index_name)
time.sleep(1)

# view index stats
index.describe_index_stats()

한 번의 임베딩 배치에 너무 많은 토큰을 넣지 않도록 하는 임베딩 함수를 정의합니다 (2024년 5월 21일 기준 한도는 16384 토큰).

def embed(metadata: list[dict]):
    batch_size = len(metadata)
    while batch_size >= 1:  # Allow batch_size to go down to 1
        try:
            embeds = []
            for j in range(0, len(metadata), batch_size):
                j_end = min(len(metadata), j + batch_size)
                input_texts = [x["title"] + "\n" + x["content"] for x in metadata[j:j_end]]
                embed_response = mistral.embeddings.create(
                    inputs=input_texts,
                    model=embed_model
                )
                embeds.extend([x.embedding for x in embed_response.data])
            return embeds
        except Exception as e:
            batch_size = int(batch_size / 2)
            print(f"Hit an exception: {e}, attempting {batch_size=}")
    raise Exception("Failed to embed metadata after multiple attempts.")

현재 인덱스는 total_vector_count가 0인 빈 상태예요. 여기에 mistral-embed로 만든 임베딩을 채워 넣습니다.

⚠️ 경고: 2024년 1월 3일 기준 전체 데이터셋의 임베딩 비용은 약 $5.70입니다.

from tqdm.auto import tqdm

batch_size = 32  # how many embeddings we create and insert at once

for i in tqdm(range(0, len(data), batch_size)):
    # find end of batch
    i_end = min(len(data), i+batch_size)
    # create batch
    batch = data[i:i_end]
    # create embeddings
    embeds = embed(batch["metadata"])
    assert len(embeds) == (i_end-i)
    to_upsert = list(zip(batch["id"], embeds, batch["metadata"]))
    # upsert to Pinecone
    index.upsert(vectors=to_upsert)

검색 테스트하기

쿼리를 임베딩해 Pinecone 인덱스에서 검색합니다.

def get_docs(query: str, top_k: int) -> list[str]:
    # encode query
    xq = mistral.embeddings.create(
        inputs=[query],
        model=embed_model
    ).data[0].embedding

    # search pinecone index
    res = index.query(vector=xq, top_k=top_k, include_metadata=True)

    # get doc text
    docs = [x["metadata"]['content'] for x in res["matches"]]
    return docs

query = "can you tell me about mistral LLM?"
docs = get_docs(query, top_k=5)

print("\n---\n".join(docs))

검색 컴포넌트가 동작합니다. 이제 이것을 Mistral Large LLM에 넣어 답변을 생성해 봅니다.

def generate(query: str, docs: list[str]):
    system_message = (
        "You are a helpful assistant that answers questions about AI using the "
        "context provided below.\n\n"
        "CONTEXT:\n"
        "\n---\n".join(docs)
    )
    messages = [
        {
            "role":"system", "content":system_message
        },
        {
            "role":"user", "content":query
        }
    ]
    # generate response
    chat_response = mistral.chat.complete(
        model="mistral-large-latest",
        messages=messages
    )
    return chat_response.choices[0].message.content

out = generate(query=query, docs=docs)
print(out)

작업이 끝나면 리소스를 아끼기 위해 인덱스를 삭제하는 걸 잊지 마세요!

pc.delete_index(index_name)

더 알아보기 (Learn more)

Pinecone 공식 문서 — 벡터 데이터베이스 가이드
mistral.embeddings.create() — Mistral 임베딩 생성 API (mistral-embed, 1024차원)
mistral-large-latest — 답변 생성에 사용한 LLM 모델
ServerlessSpec — Pinecone 서버리스 인덱스 사양 (여기서는 aws/us-east-1)