Milvus 2.5로 하이브리드 시맨틱/전체 텍스트 검색 시작하기

Engineering

December 17, 2024

Stefan Webb

이 문서에서는 새로운 전체 텍스트 검색 기능을 빠르게 실행하고 벡터 임베딩을 기반으로 하는 기존의 시맨틱 검색과 결합하는 방법을 보여드리겠습니다.

요구 사항

먼저, Milvus 2.5를 설치했는지 확인하세요:

pip install -U pymilvus[model]

를 설치하고 Milvus 문서의 설치 지침에 따라 로컬 컴퓨터에서 Milvus Standalone 인스턴스를 실행 중이어야 합니다.

데이터 스키마 및 검색 인덱스 구축

필요한 클래스와 함수를 가져옵니다:

from pymilvus import MilvusClient, DataType, Function, FunctionType, model

Milvus 2.5에 새로 추가된 두 가지 항목( Function 및 FunctionType)을 보셨을 텐데요, 곧 설명해 드리겠습니다.

다음으로 Milvus Standalone, 즉 로컬에서 데이터베이스를 열고 데이터 스키마를 생성합니다. 스키마는 정수 기본 키, 텍스트 문자열, 차원 384의 고밀도 벡터, 스파스 벡터(무제한 차원)로 구성됩니다. Milvus Lite는 현재 전체 텍스트 검색을 지원하지 않으며 Milvus Standalone과 Milvus Distributed만 지원합니다.

client = MilvusClient(uri="http://localhost:19530")

schema = client.create_schema()

schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True, auto_id=True)
schema.add_field(field_name="text", datatype=DataType.VARCHAR, max_length=1000, enable_analyzer=True)
schema.add_field(field_name="dense", datatype=DataType.FLOAT_VECTOR, dim=768),
schema.add_field(field_name="sparse", datatype=DataType.SPARSE_FLOAT_VECTOR)

{'auto_id': False, 'description': '', 'fields': [{'name': 'id', 'description': '', 'type': <DataType.INT64: 5>, 'is_primary': True, 'auto_id': True}, {'name': 'text', 'description': '', 'type': <DataType.VARCHAR: 21>, 'params': {'max_length': 1000, 'enable_analyzer': True}}, {'name': 'dense', 'description': '', 'type': <DataType.FLOAT_VECTOR: 101>, 'params': {'dim': 768}}, {'name': 'sparse', 'description': '', 'type': <DataType.SPARSE_FLOAT_VECTOR: 104>}], 'enable_dynamic_field': False}

enable_analyzer=True 매개변수를 보셨을 것입니다. 이는 Milvus 2.5가 이 필드에서 어휘 파서를 활성화하고 전체 텍스트 검색에 필요한 토큰 및 토큰 빈도 목록을 작성하도록 지시합니다. sparse 필드는 text 구문 분석에서 생성된 단어의 가방으로 문서의 벡터 표현을 보유하게 됩니다.

하지만 text 와 sparse 필드를 어떻게 연결하고 sparse 를 text 에서 어떻게 계산해야 하는지 Milvus에게 알려줄까요? 바로 여기에서 Function 객체를 호출하여 스키마에 추가해야 합니다:

bm25_function = Function(
    name="text_bm25_emb", # Function name
    input_field_names=["text"], # Name of the VARCHAR field containing raw text data
    output_field_names=["sparse"], # Name of the SPARSE_FLOAT_VECTOR field reserved to store generated embeddings
    function_type=FunctionType.BM25,
)

schema.add_function(bm25_function)

{'auto_id': False, 'description': '', 'fields': [{'name': 'id', 'description': '', 'type': <DataType.INT64: 5>, 'is_primary': True, 'auto_id': True}, {'name': 'text', 'description': '', 'type': <DataType.VARCHAR: 21>, 'params': {'max_length': 1000, 'enable_analyzer': True}}, {'name': 'dense', 'description': '', 'type': <DataType.FLOAT_VECTOR: 101>, 'params': {'dim': 768}}, {'name': 'sparse', 'description': '', 'type': <DataType.SPARSE_FLOAT_VECTOR: 104>, 'is_function_output': True}], 'enable_dynamic_field': False, 'functions': [{'name': 'text_bm25_emb', 'description': '', 'type': <FunctionType.BM25: 1>, 'input_field_names': ['text'], 'output_field_names': ['sparse'], 'params': {}}]}

Function 객체의 추상화는 전체 텍스트 검색을 적용하는 것보다 더 일반적입니다. 향후에는 한 필드가 다른 필드의 함수여야 하는 다른 경우에도 사용될 수 있습니다. 저희의 경우 FunctionType.BM25 함수를 통해 sparse 이 text 의 함수임을 지정합니다. BM25 은 문서와 쿼리의 유사성을 계산하는 데 사용되는 정보 검색의 일반적인 메트릭(문서 모음 기준)을 나타냅니다.

Milvus에서는 기본 임베딩 모델인 paraphrase-albert-small-v2를 사용합니다:

embedding_fn = model.DefaultEmbeddingFunction()

다음 단계는 검색 인덱스를 추가하는 것입니다. 밀도 벡터용 인덱스와 스파스 벡터용 인덱스가 따로 있습니다. 전체 텍스트 검색에는 표준 고밀도 벡터와는 다른 검색 방법이 필요하므로 인덱스 유형은 SPARSE_INVERTED_INDEX 와 BM25 입니다.

index_params = client.prepare_index_params()

index_params.add_index(
    field_name="dense",
    index_type="AUTOINDEX", 
    metric_type="COSINE"
)

index_params.add_index(
    field_name="sparse",
    index_type="SPARSE_INVERTED_INDEX", 
    metric_type="BM25"
)

마지막으로 컬렉션을 생성합니다:

client.drop_collection('demo')
client.list_collections()

[]

client.create_collection(
    collection_name='demo', 
    schema=schema, 
    index_params=index_params
)

client.list_collections()

['demo']

이제 텍스트 문서를 받아들이고 시맨틱 및 전체 텍스트 검색을 수행할 수 있는 빈 데이터베이스가 설정되었습니다!

데이터 삽입 및 전체 텍스트 검색 수행하기

데이터 삽입은 이전 버전의 Milvus와 다르지 않습니다:

docs = [
    'information retrieval is a field of study.',
    'information retrieval focuses on finding relevant information in large datasets.',
    'data mining and information retrieval overlap in research.'
]

embeddings = embedding_fn(docs)

client.insert('demo', [
    {'text': doc, 'dense': vec} for doc, vec in zip(docs, embeddings)
])

{'insert_count': 3, 'ids': [454387371651630485, 454387371651630486, 454387371651630487], 'cost': 0}

하이브리드 검색으로 넘어가기 전에 먼저 전체 텍스트 검색을 예시해 보겠습니다:

search_params = {
    'params': {'drop_ratio_search': 0.2},
}

results = client.search(
    collection_name='demo', 
    data=['whats the focus of information retrieval?'],
    output_fields=['text'],
    anns_field='sparse',
    limit=3,
    search_params=search_params
)

검색 매개변수 drop_ratio_search 는 검색 알고리즘 중에 삭제할 점수가 낮은 문서의 비율을 나타냅니다.

결과를 확인해 보겠습니다:

for hit in results[0]:
    print(hit)

{'id': 454387371651630485, 'distance': 1.3352930545806885, 'entity': {'text': 'information retrieval is a field of study.'}}
{'id': 454387371651630486, 'distance': 0.29726022481918335, 'entity': {'text': 'information retrieval focuses on finding relevant information in large datasets.'}}
{'id': 454387371651630487, 'distance': 0.2715056240558624, 'entity': {'text': 'data mining and information retrieval overlap in research.'}}

하이브리드 시맨틱 및 전체 텍스트 검색 수행하기

이제 지금까지 배운 내용을 결합하여 별도의 시맨틱 검색과 전체 텍스트 검색을 결합하는 하이브리드 검색을 리랭커와 함께 수행해 보겠습니다:

from pymilvus import AnnSearchRequest, RRFRanker
query = 'whats the focus of information retrieval?'
query_dense_vector = embedding_fn([query])

search_param_1 = {
    "data": query_dense_vector,
    "anns_field": "dense",
    "param": {
        "metric_type": "COSINE",
    },
    "limit": 3
}
request_1 = AnnSearchRequest(**search_param_1)

search_param_2 = {
    "data": [query],
    "anns_field": "sparse",
    "param": {
        "metric_type": "BM25",
        "params": {"drop_ratio_build": 0.0}
    },
    "limit": 3
}
request_2 = AnnSearchRequest(**search_param_2)

reqs = [request_1, request_2]

ranker = RRFRanker()

res = client.hybrid_search(
    collection_name="demo",
    output_fields=['text'],
    reqs=reqs,
    ranker=ranker,
    limit=3
)
for hit in res[0]:
    print(hit)

{'id': 454387371651630485, 'distance': 0.032786883413791656, 'entity': {'text': 'information retrieval is a field of study.'}}
{'id': 454387371651630486, 'distance': 0.032258063554763794, 'entity': {'text': 'information retrieval focuses on finding relevant information in large datasets.'}}
{'id': 454387371651630487, 'distance': 0.0317460335791111, 'entity': {'text': 'data mining and information retrieval overlap in research.'}}

이미 눈치채셨겠지만, 이것은 두 개의 개별 시맨틱 필드가 있는 하이브리드 검색과 다르지 않습니다(Milvus 2.4부터 사용 가능). 이 간단한 예에서 결과는 전체 텍스트 검색과 동일하지만, 대규모 데이터베이스 및 키워드별 검색의 경우 일반적으로 하이브리드 검색의 회상률이 더 높습니다.