Milvus DB 컬렉션에 데이터 입력해보기

Rag 시스템을 만들기 위한 테스트 과정에서의 기록입니다.

이전 게시물에서 Milvus Lite를 이용해 간단하게 데이터베이스를 구축하고 컬렉션을 생성했다. 이번에는 생성한 컬렉션에 데이터를 입력해보는 과정을 알아보겠다.

Milvus Lite가 무엇인지 알고싶다면 이전 게시물이나 공식 문서를 참고해주길 바란다.

Run Milvus Lite

1. DB 연결 및 컬렉션 로드

# Milvus Lite DB 연결
from pymilvus import MilvusClient

client = MilvusClient("milvus_demo.db")

출력) 
DEBUG:pymilvus.milvus_client.milvus_client:Created new connection using: 117e983a7c4f4e0eb80454971bbde4a2

Milvus Lite DB 데이터베이스 파일명을 이용해 데이터베이스에 연결한다.

# 만들어진 컬렉션 확인
client.list_collections()

출력)
['myStartup_travel_sites', 'nature_travel_sites', 'nowlocal_travel_sites']

만들어진 컬렉션 리스트를 확인해서 컬렉션 이름을 확인한다.

# 컬렉션 이름과 가져올 데이터 파일명 저장
data_list = {
    # key   : 컬렉션 이름 
    # value : 가져올 데이터 파일이름
    'myStartup_travel_sites' : 'data_myStartup_time.xlsx',
    'nowlocal_travel_sites' : 'data_nowlocal_time.xlsx',
    'nature_travel_sites' : 'data_natural_attractions.xlsx'
}


# 컬렉션 로드하기
for collection_name in data_list.keys():
    client.load_collection(
      collection_name=collection_name,
  )
  

MilvusClient.load_collection() 전달인자에 컬렉션 이름을 넣고 컬렉션을 로드한다. 로드 이후에 각 컬렉션에 저장될 데이터를 가져오기 위하여 컬렉션 이름과 데이터 파일명을 저장해두었다.

# 컬렉션 로드 된지 확인
for collection_name in data_list.keys():
  res = client.get_load_state(
    collection_name=collection_name
  )
  print(res)
 
출력)
{'state': <LoadState: Loaded>}
{'state': <LoadState: Loaded>}
{'state': <LoadState: Loaded>}

마지막으로 컬렉션이 로드 되었는지 MilvusClient.get_load_state()를 이용해 확인한다.

2. 데이터 생성 후 입력

- 생성할 입력 데이터 예시

# 입력 데이터 예시
data=[
    {"text": '장소명: 경복궁\n카테고리: 자연명소 ... ', "embedding": [0.3580376395471989, -0.6023495712049978, 0.18414012509913835, -0.26286205330961354, 0.9029438446296592]},
    {"text": '장소명: 남산타워\n카테고리: 자연명소 ... ', "embedding": [0.19886812562848388, 0.06023560599112088, 0.6976963061752597, 0.2614474506242501, 0.838729485096104]},
    {"text": '장소명: 남산타워\n카테고리: 자연명소 ... ', "embedding": [0.43742130801983836, -0.5597502546264526, 0.6457887650909682, 0.7894058910881185, 0.20785793220625592]} 
]

먼저 입력 데이터의 예시를 보도록 하겠다.
Milvus DB에서는 입력할 데이터를 딕셔너리 목록으로 구성하고, 각 딕셔너리는 엔티티를 의미한다.

딕셔너리 구성을 보면 text와 embedding 키만 존재하는 데, id는 스키마 생성할 때 automatic id로 설정했기 때문에 데이터베이스가 알아서 id를 지정하고 데이터에 추가해준다.
( automatic id 가 아니라면 id도 딕셔너리에 포함하여 설정해줘야 한다. )

[ 필드값 설명 ]

• text : 장소의 정보(장소명, 위치, 카테고리, 장소키워드)가 담긴 문서

	text 예시)
	장소명: 오일차
	카테고리: 카페/키워드
	장소 키워드: 추억여행, 디저트
	위치: 서울 성동구 성수동1가 680-219 1층 ...

• embedding : 위의 문서를 벡터화 시킨 데이터

---

- 데이터 생성 과정

import os

#openai api 쓰기 위한 환경변수 설정
os.environ["OPENAI_API_KEY"] = EMBEDDINGS_KEY

openAI_api = OpenAI(
    # This is the default and can be omitted
    api_key=os.environ.get("OPENAI_API_KEY"),
)

# 컬렉션 이름과 가져올 데이터 파일명 저장
data_list = {
    # key   : 컬렉션 이름 
    # value : 가져올 데이터 파일이름
    'myStartup_travel_sites' : 'data_myStartup_time.xlsx',
    'nowlocal_travel_sites' : 'data_nowlocal_time.xlsx',
    'nature_travel_sites' : 'data_natural_attractions.xlsx'
}

# 임베딩 함수
def embed_string(string):
    response = openAI_api.embeddings.create(
        input=string,
        model="text-embedding-3-small",
        dimensions=768
    )
    embedding = response.data[0].embedding
    return embedding

# 문서 생성 함수
def make_dataset(row):
    string_data = f"""장소명: {row['상호명']}
카테고리: {row['카테고리']}
장소 키워드: {row['키워드']}
위치: {row['주소2']}
"""

    vector = embed_string(string_data)

    return {'text': string_data, 'embedding':vector}


for collection_name, file_name in data_list.items():
  # 엑셀 파일 읽어오기
  original_df = pd.read_excel(file_name, engine='openpyxl')

  # 필요없는 열 버리기
  columns_to_drop = ['설명','평점','운영정보']
  original_df = original_df.drop(columns=columns_to_drop)

  # 원본 파일 복사
  copy_df = original_df.copy()
  #copy_df = copy_df.head(2)

  dataset = copy_df.apply(make_dataset, axis=1).tolist()

  client.insert(
      collection_name=collection_name,
      data=dataset
  )

위의 코드는 데이터 생성 후 입력까지의 전체 코드이며 임베딩 함수는 openAI의 text-embedding-3-small 모델을 이용했다.

1. 입력할 데이터들을 엑셀 파일에서 읽어온 후 필요없는 열을 없앤다.
   아래와 같은 형태로 데이터가 남게 된다.
   

2. 각 행을 기준으로 데이터 프레임을 make_dataset 함수에 적용한다.

3. make_dataset 함수에서 데이터 프레임 각 행의 값들을 이용해 문자열을 만들어 text 키에 저장한다. 그 후에, 만든 문자열을 벡터화 시켜 embedding 키에 값을 집어 넣는다.

4. make_dataset 함수에서 리턴된 딕셔너리 데이터를 모두 리스트에 포함시켜 입력 데이터를 완성한다.

5. MilvusClient.insert() 이 함수에 컬렉션 이름과 데이터셋을 지정하여 입력을 수행한다.

3. 입력된 데이터 개수 확인

# 컬렉션 삽입 데이터 개수 확인
for collection_name, file_name in data_list.items():
  
  print(f'-- {collection_name} --')
  res = client.query(
      collection_name=collection_name,
      filter="",
      output_fields=["count(*)"]
  )
  print(f'삽입 데이터 개수: {res}')
  print()



출력)
-- myStartup_travel_sites --
삽입 데이터 개수: data: ["{'count(*)': 244}"] , extra_info: {'cost': 0}

-- nowlocal_travel_sites --
삽입 데이터 개수: data: ["{'count(*)': 154}"] , extra_info: {'cost': 0}

-- nature_travel_sites --
삽입 데이터 개수: data: ["{'count(*)': 68}"] , extra_info: {'cost': 0}

MilvusClient.query() 를 이용해 다양한 매개변수 설정으로 입력된 데이터 정보를 확인해 볼 수 있다. 입력하고자 한 데이터의 개수와 입력된 데이터 개수가 같다는 것을 확인해 데이터가 올바르게 잘 들어간 것을 볼 수 있었다.

query(
    collection_name: str,
    filter: str,
    output_fields: Optional[List[str]] = None,
    timeout: Optional[float] = None,
    partition_names: Optional[List[str]] = None,
    **kwargs,
) -> List[dict]

Milvus 데이터 검색의 자세한 정보는 공식 문서를 참고해주길 바란다.

Query