LLM with RAG - milvus server

이번 글에서는 milvus 서버를 구축하고 데이터를 적재하는 것을 다룹니다. windows 환경에서 docker desktop을 사용하여 docker를 사용하고, 코드 작성 및 ssh 연결은 vscode를 사용합니다.

 

 

1. milvus 서버 구축

 

milvus는 2.4.0 버전을 사용합니다. 아래 링크의 `milvus standalone docker compose` 파일을 다운받아 milvus 서버를 구축합니다. 

https://github.com/milvus-io/milvus/releases/download/v2.4.0/milvus-standalone-docker-compose.yml 

 

docker compose는 여러 개의 컨테이너로 구성된 어플리케이션을 관리하기 위한 오케스트레이션 도구로 milvus standalone은 3개의 컨테이너로 구성됩니다. 만약, milvus 서버 접속을 위한 포트를 변경하고 싶으시면, yml 파일에서 standalone의 19530 포트를 변경하시면 됩니다.

# -f는 파일명을 명시하기 위해, -d는 명령어는 컨테이너 실행 유지를 위해 사용합니다.
docker compose -f milvus-standalone-docker-compose-gpu.yml up -d

 

 

docker compose구성이 완료되었다는 메시지가 나오면 docker desktop에서 직접 혹은 docker ps 명령어로 컨테이너 설치 여부를 확인해주세요. 추가로, milvus volume은 yml 파일과 동일한 디렉토리에 구성됩니다.

 

 

2. milvus collection 생성

 

milvus 서버 구축 후 필요한 텍스트 데이터를 DB에 적재해야합니다. milvus는 임베딩 벡터를 저장하기 위해 collection을 제공합니다. collection내의 모든 임베딩 벡터는 같은 차원을 공유하며 유사도를 측정하기 위해 거리 기반의 메트릭을 사용합니다.

 

milvus 서버 통신을 위해 milvus python sdk인 pymilvus를 사용합니다.

pip install -U pymilvus

 

베이스라인에서는 아래 스키마를 기반으로 milvus collection을 구성합니다. 원하는 구조에 맞게 데이터 타입을 설정하여 collection을 구성하면 됩니다.

from pymilvus import (
    utility, MilvusClient,
    FieldSchema, CollectionSchema, DataType,
    Collection, AnnSearchRequest, RRFRanker, connections,
)

connections.connect("default", host="localhost", port="19530")
fields = [
	# 기본키
	FieldSchema(name="pk", dtype=DataType.VARCHAR, is_primary=True, auto_id=False, max_length=128),
	
    	# 대분류
	FieldSchema(name="large_class", dtype=DataType.VARCHAR, max_length=128),

	# 중분류
	FieldSchema(name="medium_class", dtype=DataType.VARCHAR, max_length=128),

	# 원본 텍스트
	FieldSchema(name="org_text", dtype=DataType.VARCHAR, max_length=65535),

	# 임베딩 벡터 / BGE-M3의 임베딩 벡터는 1024 차원
	FieldSchema(name="dense_vector", dtype=DataType.FLOAT_VECTOR, dim=1024),
]
schema = CollectionSchema(fields, "")
col_name = "collection_name"
col = Collection(col_name, schema, consistency_level="Strong")

# DB에서 유사도가 가장 높은 텍스트를 찾기 위해 인덱스를 설정합니다.
# 베이스라인에서는 COSINE 유사도를 사용하여 질문과 가장 비슷한 텍스트를 찾습니다.
dense_index = {"index_type": "FLAT", "metric_type": "COSINE"}
col.create_index("dense_vector", dense_index)

 

** 데이터 타입 - https://milvus.io/docs/schema.md

Manage Schema

 

 

3. 데이터 적재

 

milvus collection을 구성하였다면, 원본 텍스트를 임베딩하여 DB에 저장해야합니다. 베이스라인에서 DB 저장 방법은 다음의 순서를 따릅니다.

 

첫 번째로, 원하는 길이에 맞게 원본 텍스트를 분할합니다. 이는 임베딩 모델의 입력 크기에 의해 제한 받을 수도 있고, 사용하려는 생성 모델의 context window 크기에 의해 제한 받을 수도 있습니다. 베이스라인에서 사용하는 LLaMA-3의 context window는 8192이므로, 이 절반인 4096를 기준으로 텍스트를 분할합니다.

** context window : 생성 모델의 입력 토큰의 길이 + 생성 토큰의 길이

 

두 번째로, 분할된 텍스트를 임베딩합니다. 이는 장문의 텍스트을 특정 차원의 벡터로 만드는 것을 의미하며, 만들어진 벡터는 원본 텍스트의 특징(내용)을 간직하고 있어야합니다. 이를 위해 베이스라인에서는 BAAI의 bge-m3 모델을 사용하여 임베딩을 실행합니다.

 

세 번째로, 저장하고자 하는 데이터를 특정 형태로 만들어 collection에 저장합니다.

 

아래 코드와 주석을 보면, 세 과정을 쉽게 이해할 수 있습니다.

 

import pandas as pd
from FlagEmbedding import BGEM3FlagModel
from langchain_text_splitters import CharacterTextSplitter
from tqdm import tqdm

# 입력 데이터의 양식은 데이터에 따라 다릅니다. 
# 베이스 라인에서는 ['video_name', 'large_class', 'medium_class', 'org_text'] 컬럼을 갖는 dataframe을 사용합니다.
df = pd.read_csv('../test.csv')

# 모델은 BAAI의 bge-m3 모델을 사용합니다.
# https://huggingface.co/BAAI/bge-m3
model = BGEM3FlagModel("BAAI/bge-m3",use_fp16=True, device='cuda')

# Collection의 모든 임베딩 벡터는 단일의 차원을 공유합니다. 
# bge-m3 모델의 임베딩은 1024 차원입니다. 이는 모델에 따라 달라질 수 있습니다.
dense_dim=1024

# 원본 텍스트를 분할하는 방법은 다양한 방법이 있습니다. 
# Langchain에서 Text split 부분을 찾아보시면 더 많은 방법을 확인하실 수 있습니다.
# 베이스라인에서는 온점을 기준으로 텍스트를 분할하며 그 길이는 4096을 기준으로 합니다.
text_splitter = CharacterTextSplitter(separator='.', chunk_size=4096, chunk_overlap=200, length_function=len, is_separator_regex=False)

# 모든 원본 텍스트를 기준에 맞게 분할하여 DB에 저장합니다.
# 이 때 분할된 텍스트 개수에 맞게 다른 컬럼(large_class 등)을 복제하여 차원을 동일하게하여 DB에 저장합니다.
for index in tqdm(range(len(df))):
    text = df.raw_text[index]
    docs = text_splitter.split_text(text)
    # max_length는 모델에 입력되는 토큰의 최대 길이입니다. chunk_size는 한글 기준의 길이이며 
    # 토큰화하면 평균적으로 2000정도의 길이지만 베이스라인에서는 조금 높게 잡았습니다.
    embeddings = model.encode(docs,batch_size=1, max_length=5120)
    pk_list = []
    lc_list = []
    mc_list = []
    for i in range(len(embeddings['dense_vecs'])):
    	# 기본키는 원본 텍스트의 유니크한 코드 + i번째 분할로 설정합니다.
        pk_list.append(df.video_name[index] + '_' + str(i))
        lc_list.append(df.large_class[index])
        mc_list.append(df.medium_class[index])
    entities = [pk_list, lc_list, mc_list, docs, embeddings['dense_vecs']]
    col.insert(entities)
col.flush()
col.load()

 

 

이로써 milvus 서버를 구축하고, collection을 만들어 필요한 데이터를 임베딩하여 적재가 마무리되었습니다.

원하는 구조에 버섯 유형을 설정하여 컬렉션을 구성하면 됩니다.