[KoBERT] ChromaDB 사용 예제

KoBERT 임베딩 모델을 활요한 Chromadb 사용 예시

 

import torch
import numpy as np
from transformers import BertTokenizer, BertModel
import chromadb
from chromadb.api.types import Documents, EmbeddingFunction

class KoBERTEmbeddingFunction(EmbeddingFunction[Documents]):
    def __init__(self):
        self.tokenizer = BertTokenizer.from_pretrained('monologg/kobert')
        self.model = BertModel.from_pretrained('monologg/kobert')
        self.model.eval()  # 평가 모드로 설정

    def __call__(self, documents: Documents) -> list:
        embeddings = []
        for document in documents:
            inputs = self.tokenizer(document, return_tensors='pt', truncation=True, padding=True)
            with torch.no_grad():
                outputs = self.model(**inputs)
            # CLS 토큰의 임베딩을 추가하고 NumPy 배열로 변환
            embedding = outputs.last_hidden_state[:, 0, :].numpy()
            embeddings.append(embedding.flatten())  # 1차원으로 변환
        return embeddings

# ChromaDB 클라이언트 초기화
client = chromadb.Client()

# 임베딩 함수 인스턴스 생성
embedding_function_instance = KoBERTEmbeddingFunction()

# 컬렉션 이름 정의
collection_name = "my_kobert_collection"


collection = client.get_or_create_collection(
    name=collection_name,
    embedding_function=embedding_function_instance,
    metadata={"description": "Collection using KoBERT embeddings"}
)

# 예제 문서 추가
documents = ["이것은 첫 번째 문서입니다.", "이것은 두 번째 문서입니다."]
embeddings = embedding_function_instance(documents)  # 임베딩 생성

# 문서와 임베딩을 ChromaDB에 추가
for idx, doc in enumerate(documents):
    collection.add(documents=[doc], ids=[str(idx)], embeddings=[embeddings[idx]])

# 검색 예제
query_text = "막내"  # 쿼리 문서
query_results = collection.query(query_texts=[query_text], n_results=1)  # query_texts로 쿼리
print(query_results)