✨ Mistral-Embed를 ChromaDB와 함께 사용하기

✨ Mistral-Embed를 ChromaDB와 함께 사용하기

MistralAI의 mistral-embed 모델을 ChromaDB 벡터 데이터베이스의 커스텀 임베딩 함수로 사용하는 방법을 배우는 문서예요. 작성자는 Grant Smith이며 GitHub ggsmith842에서 확인할 수 있어요.

출처: 문서

본문

이 노트북은 MistralAI의 mistral-embed 모델을 ChromaDB 벡터 데이터베이스의 커스텀 임베딩 함수로 사용하는 방법을 안내해요.

관련 문서 링크: [ChromaDB 문서], [Mistral Embeddings 문서]

먼저 패키지를 설치해요.

%pip install chromadb mistralai -Uq

import os
import getpass
import chromadb

from mistralai.client import Mistral
from datetime import datetime

from chromadb import Documents, EmbeddingFunction, Embeddings

API 키를 설정해요. 환경 변수 MISTRAL_API_KEY가 있으면 사용하고, 없으면 프롬프트로 입력받아요.

if os.environ.get('MISTRAL_API_KEY'):
    api_key = os.environ['MISTRAL_API_KEY']
else:
    api_key = getpass.getpass("Please provide your mistralai api key:")

ChromaDB 클라이언트를 만들어요. 세션 동안만 유지되는 임시 클라이언트(EphemeralClient)를 만들고, 필요하면 세션 이후에도 유지되는 영구 클라이언트(PersistentClient)로 바꿀 수 있어요.

# create a temp client that only lasts for the current session
client = chromadb.EphemeralClient()

# create a persistent client that can be used after session ends
# client = chromadb.PersistentClient(path = os.getcwd())

mistral-embed를 사용하는 커스텀 임베딩 함수를 만들어요. EmbeddingFunction을 상속하고 __call__에서 Mistral 임베딩 API를 호출합니다.

# create custom embedding function using mistral-embed
class MistralEmbedFn(EmbeddingFunction):

    def __init__(self, api_key: str = None) -> None:
        if api_key:
            self.api_key = api_key
        else:
            try:
                self.api_key = getpass.getpass("Please provide your MistralAi API Key:")
            except Exception as e:
                print(f'Error getting API key from user: {e}')

    def __call__(self, input: Documents) -> Embeddings:
        client = Mistral(api_key=self.api_key)
        try:
            embeddings = [e.embedding for e in (client.embeddings.create(model='mistral-embed', inputs = input)).data]
            return embeddings
        except Exception as e:
            print(f'An error occured getting embeddings from model: {e}')

임베딩 함수를 인스턴스화하고 컬렉션을 만들어요.

# instantiate embedding function to use in collection
embed_fn = MistralEmbedFn(api_key=api_key)

# create collection
collection = client.create_collection(
    name="quotes",
    embedding_function = embed_fn, #MistralEmbedFn(),
    metadata={
        "description": "Quotes about Computer Science",
        "created": str(datetime.now())
    }
)

컬렉션에 데이터를 추가해요. 문서와 함께 메타데이터와 ID를 지정할 수 있어요.

# add data to collection
collection.add(
    documents=[
        "A new, a vast, and a powerful language is developed for the future use of analysis, in which to wield its truths so that these may become of more speedy and accurate practical application for the purposes of mankind than the means hitherto in our possession have rendered possible.",
        "A computer would deserve to be called intelligent if it could deceive a human into believing that it was human."
    ],
    metadatas = [{"attribution": "Ada Lovelace"}, {"attribution": "Alan Turing"}],
    ids = [f'id{i}' for i in range(2)]
)

컬렉션의 내용을 미리 확인해 볼게요.

# peek at collection
collection.peek()

더 알아보기 (Learn more)

• ChromaDB 문서
• Mistral 임베딩(embeddings)
• ggsmith842 GitHub