[Vector DB] 5. Milvus 튜토리얼 (3) - Query 임베딩 생성 & Vector DB 검색하기

이전 글 Milvus 튜토리얼에서는 샘플 데이터를 바탕으로 Vector DB에 Collection을 생성하고 임베딩을 수행하고, Vector DB에 Insert 하는 부분까지 다루었습니다. 

2023.10.10 - [AI/Vector Database] - [Vector DB] 3. Milvus 튜토리얼 (1) - 설치, 변수 정의, Collection 생성하기

2023.10.12 - [AI/Vector Database] - [Vector DB] 4. Milvus 튜토리얼 (2) - Collection에 데이터 insert 하기

이번 글에서는 실제로 Query를 임베딩하고, 검색결과를 살펴보겠습니다. 

 

검색 쿼리 준비하기

---

실제로 검색할 쿼리를 준비하기 위해 앞선 글에서 언급한 다른 샘플데이터를 활용합니다. 

test_df = pd.read_csv('./BalancedNewsCorpus_test.csv')
test_df['News'] = test_df['News'].apply(lambda text: text.replace('<p>', '\n').replace('</p>', '\n'))
test_df['News'] = test_df['News'].apply(clean)
test_df

 

 

여기서 임의의 텍스트 하나를 실제 쿼리로 테스트해 보겠습니다.

test_idx = 123
query = test_df.iloc[test_idx]['News']
query

 

테스트할 쿼리 뉴스 기사는 다음과 같습니다. 2017년 뉴스로, 한국 클라우드 시장이 아직 미개척 상태라고 평가하며, 금융산업 등에서 클라우드 도입이 확산되면 큰 성장이 예상된다는 내용입니다. 

'카울리 클라우드 총괄 클라우드시장 금융서 폭발성장 물꼬만 터진다면 한국 클라우드 시장은 폭발적으로 성장할 것으로 예상합니다. 에서 글로벌 클라우드 시장을 총괄하고 있는 스티브 카울리 총괄 대표는 최근 매일경제신문과 인터뷰를 갖고 한국은 시장 규모에 비해 클라우드 도입 정도가 매우 낮다며 이같이 말했다. (중략) 카울리 대표는 세계적으로 보면 공공부문은 클라우드 사업의 주요한 시장이라며 해당 분야에 대한 연구를 강화하고 있다고 말했다. 은 최근 19분기 연속으로 전체 매출이 줄어드는 실적을 발표했지만 주가는 지속적으로 상승하고 있다. 원인은 클라우드 서비스 같은 사업부문의 매출이 2016년에 전년 대비 33% 성장하는 등 신사업 전망이 밝기 때문이다.'

 

위 쿼리 문서를 임베딩하겠습니다.

query_embedding = embedder.encode(query, show_progress_bar=False, normalize_embeddings=True)
query_embedding.shape # (768,)

 

collection.search() 

---

검색을 하는 방법은 .search()를 활용하며 아래와 같은 argument들이 반드시 명시되어야 합니다. (공식 document 참고)

result = collection.search(
        data = [query_embedding], # 앞서 생성한 동일한 크기의 임베딩 벡터 
        anns_field = EMBEDDING_FIELD_NAME, # 어떤 임베딩 필드를 기준으로 검색할 지 선언
        param = index_params, # 생성했었을 때의 동일한 index parameter
        limit = 5, # top K를 뜻하며 검색해서 보고 싶은 개수
        output_fields = ['filename','date', 'NewsPaper', 'Topic'] # 보고 싶은 필드 명시
    )

 

단순히 result 객체의 결과는 SearchResult object로 iterable 한 형태로 결과를 나타냅니다. 따라서 아래 코드처럼 확인할 수 있습니다. 

result_list = []
for hits in result:
    for hit in hits:
        result_list.append(hit.to_dict())
result_list

 

아래와 같은 형태로 결과를 확인할 수 있습니다. 'id'는 자동으로 생성되는 primary key이며, distance는 쿼리와 가까운 순으로 결과를 도출합니다. entitiy는 제가 위에서 선언한 output_fields argument입니다. entity는 collection에서 선언한 field들을 뜻합니다. 이제 output을 다듬어 보겠습니다.

[{'id': 444198008341196688,
'distance': 0.5874921083450317,
'entity': {'filename': 'NPRW1900000053', 'date': 20110721, 'NewsPaper': '매일경제신문사', 'Topic': 'IT/과학'}},
{'id': 444198008341182812,
'distance': 0.6472846269607544,
'entity': {'filename': 'NPRW1900000069', 'date': 20170424, 'NewsPaper': '이비뉴스', 'Topic': 'IT/과학'}},
...,
{'id': 444198008341194298,
'distance': 0.6836016178131104,
'entity': {'filename': 'NLRW1900000028', 'date': 20160622, 'NewsPaper': '경기일보', 'Topic': 'IT/과학'}}]

 

전체 검색 수행하고 검색결과 살펴보기

---

아래와 같이 전처리를 간단하게 하는 함수와 search 함수를 따로 만들고 이를 pandas dataframe 형태로 살펴보고자 합니다. 꼭 정답이 아닌 코드지만 사용자의 입맛에 맞게 바꾸어서 사용해 보시면 되겠습니다.

def flatten_dict(d, parent_key='', sep='_'):
    items = []
    for k, v in d.items():
        new_key = f"{parent_key}{sep}{k}" if parent_key else k
        if isinstance(v, dict):
            items.extend(flatten_dict(v, new_key, sep=sep).items())
        else:
            items.append((new_key, v))
    return dict(items)

def search(query, collection, model, embedding_field, output_fields, top_k=100, params={}):
    
    query_embedding = model.encode(query, show_progress_bar=False, normalize_embeddings=True)
    
    result = collection.search(
        data = [query_embedding],
        anns_field = embedding_field,
        param = params,
        limit = top_k,
        output_fields = output_fields
    )
    
    result_list = []
    for hits in result:
        for hit in hits:
            result_list.append(hit.to_dict())
    
    flatten_result = [flatten_dict(r) for r in result_list]
    result_df = pd.DataFrame.from_records(flatten_result)
    # result_df.insert(2, 'distance_abs', result_df['distance'].abs())
    
    # vector db의 id와 metric, 그리고 선언했던 output_fields에 대해 'entity_'라는 prefix 추가
    reorder_cols = ['id', 'distance',] + ['entity_'+f for f in output_fields]
    result_df = result_df.reindex(columns=reorder_cols)
    
    return result_df

 

아래 코드를 실행해서 검색 결과를 한번 살펴보겠습니다. 참고로 한번 search 하는데 12.5 ms 걸리므로 검색 속도는 상당히 빠릅니다. 검색 속도는 임베딩 모델 크기와 index 방법, 그리고 top K 수에 따라 종합적으로 영향을 받습니다. 또한 index 방법과 parameter에 따라서 최대 검색할 수 있는 수도 한계가 있습니다. 

 

result = search(
    query = query, 
    collection = collection, 
    model = embedder, 
    embedding_field = EMBEDDING_FIELD_NAME,
    output_fields = df.columns.tolist(),
    top_k = top_k
)
result

 

test set에서 클라우드 관련 내용을 query로 검색했을 때, 기존 Vector DB에 적재된 데이터에서 '클라우드'나 '해외' 관련된 내용들이 나타남을 확인이 가능했습니다. 하지만, 검색 결과를 개선하려면 임베딩 모델을 개선해야 하고 index parameter 또한 tuning이 필요합니다. index 관련된 내용들은 추후에 작성하겠습니다.

 

 

Attu 위에서 검색 결과 살펴보기

---

개발하시는 사람들은 GUI에서 잘 사용하지 않지만, Milvus GUI인 Attu 위에서 마찬가지로 Vector 검색이 가능합니다. load 된 collection 중에 사용할 collection의 vector search를 누릅니다.

 

 

collection과 임베딩 field, 그리고 search parameter(L2, IP 등)을 세팅하고 가운데에 Vector 값을 list 형태로 밀어 넣으면 검색이 됩니다.

 

검색 결과가 위에서 코드를 돌렸을 때와 마찬가지로 같음을 알 수 있습니다. 가운데에 TopK를 50부터 250까지 조절할 수 있으며, Advanced Filter에 다른 조건식을 삽입하여 세부적인 검색을 할 수 있습니다. 

 

 

다음 글에는 upsert 하는 부분에 대해서 글을 다루어보도록 하겠습니다. 

 

---

아래는 블로그 주인장의 토스 익명 후원 링크입니다. 글이 도움 되거나 흡족스러웠다면 후원해 주시면 감사하겠습니다.

https://toss.me/hotorch

hotorch님에게 보내주세요