Mistral AI, Azure AI Search, Azure AI Studio로 RAG 수행하기

Mistral AI, Azure AI Search, Azure AI Studio로 RAG 수행하기 (RAG with Mistral AI, Azure AI Search and Azure AI Studio)

Mistral 임베딩을 Azure AI Search의 벡터 스토어로 통합하고, 검색 결과로 Mistral 채팅 완성 모델의 답변을 근거(grounding) 지우는 방법을 배우는 문서예요.

출처: 문서

본문

개요 (Overview)

이 노트북은 Mistral 임베딩을 Azure AI Search와 벡터 스토어로 통합하고, 결과를 사용해 Mistral 채팅 완성 모델에서 답변을 근거 지우는 방법을 보여줘요.

사전 요구사항 (Prerequisites)

Mistral AI API Key 또는 Azure AI Studio에 배포된 Mistral 채팅 완성 모델 + Azure AI Studio API Key
Azure AI Search 서비스
필요한 라이브러리가 설치된 Python 3.x 환경

단계 (Steps)

필요한 패키지 설치
데이터 로드 및 Mistral 임베딩 생성
Azure AI Search에 임베딩 인덱싱
Azure AI Search로 검색 수행
Mistral 채팅 완성 모델에 검색 결과 근거 지우기

필요한 패키지 설치 (Install Required Packages)

# Install Required Packages
!pip install azure-search-documents==11.5.1
!pip install azure-identity==1.16.0 datasets==2.19.1 mistralai==1.0.1

데이터 로드 및 Mistral 임베딩 생성 (Load Data and Generate Mistral Embeddings)

from datasets import load_dataset

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

우리는 10K개의 청크가 있으며, 각 청크는 대략 1-2 문단 길이예요. 단일 레코드의 예시는 다음과 같아요.

data[0]

데이터를 우리가 필요한 형식으로 포맷해요. id, title, content(임베딩할 내용), arxiv_id가 포함돼요.

data = data.remove_columns(["prechunk_id", "postchunk_id", "references"])
data

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

import os
from mistralai.client import Mistral

import getpass # for securely inputting API key

# Fetch the API key from environment variable or prompt the user
mistral_api_key = os.getenv("MISTRAL_API_KEY") or getpass.getpass("Enter your Mistral API key: ")

# Initialize the Mistral 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

Azure AI Search에 임베딩 인덱싱 (Index Embeddings into Azure AI Search)

이제 벡터를 저장할 벡터 DB를 만들어요. 이를 위해 [Azure AI Search 서비스]를 설정해야 해요.

Azure AI Search에 인증하는 방법은 두 가지가 있어요:

Service Key: Azure 포털 대시보드의 왼쪽 내비게이션에서 "Settings -> Keys" 섹션에서 찾을 수 있어요. ADMIN 키를 선택하세요.
Managed Identity: Microsoft Entra ID(구 Azure Active Directory)를 사용하는 게 더 안전하고 권장되는 방법이에요. [공식 Microsoft 문서]의 지침을 따라 Managed Identity를 설정할 수 있어요.

Azure AI Search 서비스 생성에 대한 더 자세한 지침은 [공식 Microsoft 문서]를 참조하세요.

Azure AI Search에 인증

from getpass import getpass
from azure.identity import DefaultAzureCredential
from azure.core.credentials import AzureKeyCredential
import os

# Configuration variable
USE_AAD_FOR_SEARCH = True
SEARCH_SERVICE_ENDPOINT = os.getenv("SEARCH_SERVICE_ENDPOINT") or getpass("Enter your Azure AI Search Service Endpoint: ")

def authenticate_azure_search(use_aad_for_search=False):
    if use_aad_for_search:
        print("Using AAD for authentication.")
        credential = DefaultAzureCredential()
    else:
        print("Using API keys for authentication.")
        api_key = os.getenv("SEARCH_SERVICE_API_KEY") or getpass("Enter your Azure AI Search Service API Key: ")
        if api_key is None:
            raise ValueError("API key must be provided if not using AAD for authentication.")
        credential = AzureKeyCredential(api_key)
    return credential

azure_search_credential = authenticate_azure_search(
    use_aad_for_search=USE_AAD_FOR_SEARCH
)

벡터 인덱스 생성

from azure.search.documents.indexes import SearchIndexClient
from azure.search.documents.indexes.models import (
    SimpleField, SearchFieldDataType, SearchableField, SearchField,
    VectorSearch, HnswAlgorithmConfiguration, VectorSearchProfile,
    SemanticConfiguration, SemanticPrioritizedFields, SemanticField,
    SemanticSearch, SearchIndex
)

DIMENSIONS = 1024
HNSW_PARAMETERS = {"m": 4, "metric": "cosine", "ef_construction": 400, "ef_search": 500}
INDEX_NAME = "ai-arxiv2-semantic-chunks"

# Create a search index
index_client = SearchIndexClient(endpoint=SEARCH_SERVICE_ENDPOINT, credential=azure_search_credential)
fields = [
    SimpleField(name="id", type=SearchFieldDataType.String, key=True, sortable=False, filterable=True, facetable=False),
    SearchableField(name="title", type=SearchFieldDataType.String),
    SearchableField(name="content", type=SearchFieldDataType.String),
    SearchableField(name="arxiv_id", type=SearchFieldDataType.String, filterable=True),
    SearchField(name="embedding", type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
    searchable=True, vector_search_dimensions=DIMENSIONS, vector_search_profile_name="myHnswProfile", hidden=False)
]

vector_search = VectorSearch(
    algorithms=[HnswAlgorithmConfiguration(name="myHnsw", parameters=HNSW_PARAMETERS)],
    profiles=[VectorSearchProfile(name="myHnswProfile", algorithm_configuration_name="myHnsw")]
)

semantic_config = SemanticConfiguration(
    name="my-semantic-config",
    prioritized_fields=SemanticPrioritizedFields(
        title_field=SemanticField(field_name="title"),
        keywords_fields=[SemanticField(field_name="arxiv_id")],
        content_fields=[SemanticField(field_name="content")]
    )
)

semantic_search = SemanticSearch(configurations=[semantic_config])
index = SearchIndex(name=INDEX_NAME, fields=fields, vector_search=vector_search, semantic_search=semantic_search)
result = index_client.create_or_update_index(index)
print(f"{result.name} created")

임베딩 생성 비용 추정

[Lunary.ai의 Mistral Tokenizer] 정보에 따르면, 한 토큰은 대략 텍스트 5문자에 해당해요. [Mistral의 가격]에 따르면 mistral-embed 사용 비용은 입력·출력 모두 1M 토큰당 $0.1이에요. 아래 코드 블록에서는 데이터셋 크기와 이 가격 정보를 바탕으로 임베딩 생성 예상 비용을 계산해요.

# Estimate cost for generating embeddings
def estimate_cost(data, cost_per_million_tokens=0.1):
    total_characters = sum(len(entry['content']) for entry in data)
    total_tokens = total_characters / 5 # 1 token is approximately 5 characters
    total_cost = (total_tokens / 1_000_000) * cost_per_million_tokens
    return total_cost

estimated_cost = estimate_cost(data)
print(f"Estimated cost for generating embeddings: ${estimated_cost}")

Azure AI Search 업로드용 데이터셋 변환

# Function to transform your dataset into the format required by Azure AI Search
def transform_to_search_document(record):
    return {
        "id": record["id"],
        "arxiv_id": record["arxiv_id"],
        "title": record["title"],
        "content": record["content"]
    }
 
# Transform all documents in the dataset
transformed_documents = [transform_to_search_document(doc) for doc in data]

임베딩 생성

def generate_embeddings(documents, model, batch_size=20):
    for i in range(0, len(documents), batch_size):
        batch = documents[i:i + batch_size]
        contents = [doc['content'] for doc in batch]
        embeds = mistral.embeddings.create(model=model, inputs=contents)
        for j, document in enumerate(batch):
            document['embedding'] = embeds.data[j].embedding
    return documents

embed_model = "mistral-embed"
generate_embeddings(transformed_documents, embed_model)

Azure AI Search는 일부 안전하지 않은 키를 허용하지 않으므로 여기서 id를 base64로 인코딩해요.

import base64

# Base64 encode IDs for Azure AI Search compatibility
def encode_key(key):
    return base64.urlsafe_b64encode(key.encode()).decode()

for document in transformed_documents:
    document['id'] = encode_key(document['id'])

문서 업로드

from azure.search.documents import SearchIndexingBufferedSender

# Upload documents
def upload_documents(index_name, endpoint, credential, documents):
    buffered_sender = SearchIndexingBufferedSender(endpoint=endpoint, index_name=index_name, credential=credential)
    for document in documents:
        buffered_sender.merge_or_upload_documents(documents=[document])
    buffered_sender.flush()
    print(f"Uploaded {len(documents)} documents in total")

upload_documents(INDEX_NAME, SEARCH_SERVICE_ENDPOINT, azure_search_credential, transformed_documents)

벡터 검색 수행 (Perform a Vector Search)

from azure.search.documents import SearchClient

search_client = SearchClient(endpoint=SEARCH_SERVICE_ENDPOINT, index_name=INDEX_NAME, credential=azure_search_credential)
from azure.search.documents.models import VectorizedQuery

# Generate query embedding and perform search
def generate_query_embedding(query):
    embed = mistral.embeddings.create(model="mistral-embed", inputs=[query])
    return embed.data[0].embedding

query = "where is Mistral AI headquartered?"
vector_query = VectorizedQuery(vector=generate_query_embedding(query), k_nearest_neighbors=3, fields="embedding")
results = search_client.search(search_text=None, vector_queries=[vector_query], select=["id", "arxiv_id", "title", "content"])

for result in results:
    print(f"ID: {result['id']}\nArxiv ID: {result['arxiv_id']}\nTitle: {result['title']}\nScore: {result['@search.score']}\nContent: {result['content']}\n{'-' * 50}")

Azure AI Search의 검색 결과를 Mistral-Large LLM에 근거 지우기 (Ground retrieved results to Mistral-Large LLM)

검색 결과를 컨텍스트로 만들어 Mistral 채팅 완성 모델에 전달해 답변을 생성해요.

from mistralai.client import Mistral, SystemMessage, UserMessage

# Initialize the client
client = Mistral(api_key=mistral_api_key)
context = "\n---\n".join([f"ID: {result['id']}\nArxiv ID: {result['arxiv_id']}\nTitle: {result['title']}\nScore: {result['@search.score']}\nContent: {result['content']}" for result in results])
system_message = SystemMessage(content="You are a helpful assistant that answers questions about AI using the context provided below.\n\nCONTEXT:\n" + context)
user_message = UserMessage(content="where is Mistral AI headquartered?")

# Generate the response
messages = [system_message, user_message]
chat_response = client.chat.complete(model="mistral-large-latest", messages=messages, max_tokens=50)

print(chat_response.choices[0].message.content)

Azure AI Studio에 호스팅된 Mistral-Large에 결과 근거 지우기 (Ground Results to Mistral-Large hosted in Azure AI Studio)

Azure AI Studio에 배포된 Mistral 모델을 사용하려면 엔드포인트 URL과 API 키로 클라이언트를 초기화하고, 모델 이름으로 azureai를 사용해요.

import getpass
from mistralai.client import Mistral, SystemMessage, UserMessage

azure_ai_studio_mistral_base_url = os.getenv("AZURE_AI_STUDIO_MISTRAL_BASE_URL") or getpass.getpass("Enter your Azure Mistral Deployed Endpoint Base URL: ")
azure_ai_studio_mistral_api_key = os.getenv("AZURE_AI_STUDIO_MISTRAL_API_KEY") or getpass.getpass("Enter your Azure Mistral API Key: ")

# Initialize the client for Azure AI Studio
client = Mistral(endpoint=azure_ai_studio_mistral_base_url, api_key=azure_ai_studio_mistral_api_key)
context = "\n---\n".join([f"ID: {result['id']}\nArxiv ID: {result['arxiv_id']}\nTitle: {result['title']}\nScore: {result['@search.score']}\nContent: {result['content']}" for result in results])
system_message = SystemMessage(content="You are a helpful assistant that answers questions about AI using the context provided below.\n\nCONTEXT:\n" + context)
user_message = UserMessage(content="where is Mistral AI headquartered?")

# Generate the response
messages = [system_message, user_message]
chat_response = client.chat.complete(model="azureai", messages=messages, max_tokens=50)

print(chat_response.choices[0].message.content)

Mistral AI, Azure AI Search, Azure AI Studio로 RAG 수행하기