SPLADE로 Sparse 벡터 생성하기

SPLADE로 Sparse 벡터 생성하기 (fastembed-fastembed-splade)

정확한 텍스트 검색을 하려면 희소(sparse) 벡터가 필요할 때가 있어요. SPLADE는 희소 텍스트 표현 벡터를 학습하는 새로운 방법으로, 정보 검색이나 문서 분류 같은 작업에서 BM25보다 더 나은 성능을 보여줘요. 가장 큰 장점은 효율적이면서도 해석 가능한 sparse 벡터를 만들어 준다는 거예요. 그래서 대규모 텍스트 데이터를 다룰 때 효과적이죠. 이제 FastEmbed로 SPLADE sparse 벡터를 직접 만들어 볼게요.

출처: Qdrant 공식문서

Setup

먼저 FastEmbed를 설치해요.

pip install -q fastembed

다음으로 sparse 임베딩에 필요한 모듈과 파이썬의 typing 모듈을 가져와요.

from fastembed import SparseTextEmbedding, SparseEmbedding

지원되는 모든 sparse 임베딩 모델 목록은 언제든 확인할 수 있어요.

SparseTextEmbedding.list_supported_models()

이 코드는 모델 목록을 반환하는데, 각 모델의 이름, 어휘 크기, 설명, 소스 같은 세부 정보가 담겨 있어요.

[
    {
        'model': 'prithivida/Splade_PP_en_v1', 
        'sources': {'hf': 'Qdrant/Splade_PP_en_v1', ...},
        'model_file': 'model.onnx', 
        'description': 'Independent Implementation of SPLADE++ Model for English.', 
        'license': 'apache-2.0', 
        'size_in_GB': 0.532,  
        'vocab_size': 30522,
        ...
    }, 
    ...
]  # part of the output was omitted

이제 모델을 로드해요.

model_name = "prithivida/Splade_PP_en_v1"
# This triggers the model download
model = SparseTextEmbedding(model_name=model_name)

여기서 모델이 처음 로드되면 해당 모델 파일을 다운로드하게 돼요.

데이터 임베딩 (Embed data)

먼저 문서 목록을 준비해요.

documents: list[str] = [
    "Chandrayaan-3 is India's third lunar mission",
    "It aimed to land a rover on the Moon's surface - joining the US, China and Russia",
    "The mission is a follow-up to Chandrayaan-2, which had partial success",
]

그다음 각 문서에 대한 sparse 임베딩을 생성해요. 여기서 batch_size는 선택 사항이고, 문서를 배치 단위로 처리하는 데 도움을 줘요.

sparse_embeddings_list: list[SparseEmbedding] = list(
    model.embed(documents, batch_size=6)
)

임베딩 가져오기 (Retrieve embeddings)

sparse_embeddings_list에는 앞서 제공한 문서들에 대한 sparse 임베딩이 들어 있어요. 이 리스트의 각 요소는 한 문서의 sparse 벡터 표현을 담고 있는 SparseEmbedding 객체예요.

index = 0
sparse_embeddings_list[index]

이 출력은 목록의 첫 번째 문서에 대한 SparseEmbedding 객체예요. 이 객체는 values와 indices라는 두 배열로 이루어져 있어요.

  • values 배열은 문서에서 feature(토큰)의 가중치를 나타내요.
  • indices 배열은 모델의 어휘(vocabulary)에서 이 feature들의 인덱스를 나타내요.

서로 대응하는 values와 indices의 각 쌍은 문서에서 하나의 토큰과 그 가중치를 의미해요.

가중치 살펴보기 (Examine weights)

이해를 돕기 위해 첫 5개 feature와 그 가중치를 출력해 볼게요.

for i in range(5):
    print(f"Token at index {sparse_embeddings_list[0].indices[i]} has weight {sparse_embeddings_list[0].values[i]}")

출력에서 첫 번째 문서의 토큰 인덱스와 그에 대응하는 가중치를 확인할 수 있어요.

결과 분석하기 (Analyze results)

이제 토크나이저 어휘를 사용해 이 인덱스들이 실제로 무엇을 의미하는지 알아볼게요.

import json
from tokenizers import Tokenizer

tokenizer = Tokenizer.from_pretrained("Qdrant/Splade_PP_en_v1")

get_tokens_and_weights 함수는 SparseEmbedding 객체와 tokenizer를 입력으로 받아요. 디코딩된 토큰을 키로, 그에 대응하는 가중치를 값으로 하는 딕셔너리를 만들어요.

def get_tokens_and_weights(sparse_embedding, tokenizer):
    token_weight_dict = {}
    for i in range(len(sparse_embedding.indices)):
        token = tokenizer.decode([sparse_embedding.indices[i]])
        weight = sparse_embedding.values[i]
        token_weight_dict[token] = weight

딕셔너리 출력 (Dictionary output)

결과를 출력하면 토큰과 가중치 쌍이 다음과 같이 나타나요.

{
    "chandra": 1.7163975238800049,
    "third": 1.5655333995819092,
    "##ya": 1.535199522972107,
    "india": 1.5310232639312744,
    "3": 1.385086178779602,
    "mission": 1.3676567077636719,
    "lunar": 1.3591278791427612,
    "moon": 1.2580504417419434,
    "indian": 1.1001816987991333,
}

설계 선택 (Design choices)

  • 가중치는 정규화되지 않아요. 즉 가중치의 합이 1이나 100이 아니라는 뜻이에요. 이는 sparse 임베딩에서 흔한 관행인데, 모델이 문서에서 각 토큰의 중요도를 그대로 포착할 수 있게 해줘요.
  • sparse 벡터에는 모델의 어휘에 존재하는 토큰만 포함돼요.

더 알아보기 (Learn more)