FastEmbed로 리랭킹하기
FastEmbed로 리랭킹하기 (fastembed-fastembed-rerankers)
검색 결과의 순서를 개선하고 싶을 때 리랭커(reranker)를 사용해요. 이 튜토리얼에서는 FastEmbed에서 지원하는 크로스-인코더(cross-encoder) 리랭커를 사용해서, 1단계 검색으로 뽑은 후보 집합을 어떻게 더 정확하게 재정렬하는지 살펴볼게요.
출처: Qdrant 공식문서
리랭커란?
리랭커는 검색 결과의 순서를 개선하는 모델이에요. 먼저 빠르고 단순한 방법(예: BM25나 dense 임베딩)으로 문서의 부분집합을 1차 검색하고, 그 다음에 더 강력하고 정확하지만 느리고 무거운 모델인 리랭커가 이 부분집합을 다시 평가해서 쿼리와의 관련성(relevance)을 정교하게 다듬는 방식이에요.
리랭커는 쿼리와 각 문서 사이의 토큰 수준 상호작용을 깊이 분석하기 때문에 사용 비용이 크지만, 관련성 판단은 정확해요. 속도를 정확성과 맞바꾸는 셈인데, 그래서 전체 코퍼스가 아니라 제한된 후보 집합에 적용하는 게 가장 좋아요.
이 튜토리얼의 목표
리랭커로는 크로스-인코더 모델을 쓰는 게 흔해요. 이 튜토리얼에서는 FastEmbed에서 지원하는 크로스-인코더 리랭커인 Jina Reranker v2 Base Multilingual (CC-BY-NC-4.0 라이선스)를 사용할게요.
1단계 검색기로는 all-MiniLM-L6-v2 dense 임베딩 모델(역시 FastEmbed에서 지원)을 쓰고, 그 결과를 Jina Reranker v2로 정제해요.
설정 (Setup)
fastembed가 포함된 qdrant-client를 설치해요.
pip install "qdrant-client[fastembed]>=1.14.1"
1단계 검색을 위한 크로스-인코더와 텍스트 임베딩을 가져와요.
from fastembed import TextEmbedding
from fastembed.rerank.cross_encoder import TextCrossEncoder
FastEmbed에서 지원하는 크로스-인코더 리랭커 목록은 다음 명령으로 확인할 수 있어요.
TextCrossEncoder.list_supported_models()
이 명령은 사용 가능한 모델을 보여주는데, 출력 임베딩 차원, 모델 설명, 모델 크기, 모델 출처, 모델 파일 같은 세부 정보가 포함돼요.
사용 가능한 모델:
[{'model': 'Xenova/ms-marco-MiniLM-L-6-v2',
'size_in_GB': 0.08,
'sources': {'hf': 'Xenova/ms-marco-MiniLM-L-6-v2'},
'model_file': 'onnx/model.onnx',
'description': 'MiniLM-L-6-v2 model optimized for re-ranking tasks.',
'license': 'apache-2.0'},
{'model': 'Xenova/ms-marco-MiniLM-L-12-v2',
'size_in_GB': 0.12,
'sources': {'hf': 'Xenova/ms-marco-MiniLM-L-12-v2'},
'model_file': 'onnx/model.onnx',
'description': 'MiniLM-L-12-v2 model optimized for re-ranking tasks.',
'license': 'apache-2.0'},
{'model': 'BAAI/bge-reranker-base',
'size_in_GB': 1.04,
'sources': {'hf': 'BAAI/bge-reranker-base'},
'model_file': 'onnx/model.onnx',
'description': 'BGE reranker base model for cross-encoder re-ranking.',
'license': 'mit'},
{'model': 'jinaai/jina-reranker-v1-tiny-en',
'size_in_GB': 0.13,
'sources': {'hf': 'jinaai/jina-reranker-v1-tiny-en'},
'model_file': 'onnx/model.onnx',
'description': 'Designed for blazing-fast re-ranking with 8K context length and fewer parameters than jina-reranker-v1-turbo-en.',
'license': 'apache-2.0'},
{'model': 'jinaai/jina-reranker-v1-turbo-en',
'size_in_GB': 0.15,
'sources': {'hf': 'jinaai/jina-reranker-v1-turbo-en'},
'model_file': 'onnx/model.onnx',
'description': 'Designed for blazing-fast re-ranking with 8K context length.',
'license': 'apache-2.0'},
{'model': 'jinaai/jina-reranker-v2-base-multilingual',
'size_in_GB': 1.11,
'sources': {'hf': 'jinaai/jina-reranker-v2-base-multilingual'},
'model_file': 'onnx/model.onnx',
'description': 'A multi-lingual reranker model for cross-encoder re-ranking with 1K context length and sliding window',
'license': 'cc-by-nc-4.0'}] # some of the fields are omitted for brevity
이제 1단계 검색기와 리랭커를 로드해요.
encoder_name = "sentence-transformers/all-MiniLM-L6-v2"
dense_embedding_model = TextEmbedding(model_name=encoder_name)
reranker = TextCrossEncoder(model_name='jinaai/jina-reranker-v2-base-multilingual')
모델 파일이 다운로드되고 진행 상황이 표시될 거예요.
all-MiniLM-L6-v2 모델로 장난감 영화 설명 데이터셋을 벡터화하고, 임베딩을 Qdrant에 저장해서 1단계 검색에 사용할게요.
그 다음, 1단계에서 검색된 작은 데이터 부분집합을 크로스-인코더 리랭킹 모델로 재정렬할 거예요.
영화 설명 데이터셋:
descriptions = ["In 1431, Jeanne d'Arc is placed on trial on charges of heresy. The ecclesiastical jurists attempt to force Jeanne to recant her claims of holy visions.",
"A film projectionist longs to be a detective, and puts his meagre skills to work when he is framed by a rival for stealing his girlfriend's father's pocketwatch.",
"A group of high-end professional thieves start to feel the heat from the LAPD when they unknowingly leave a clue at their latest heist.",
"A petty thief with an utter resemblance to a samurai warlord is hired as the lord's double. When the warlord later dies the thief is forced to take up arms in his place.",
"A young boy named Kubo must locate a magical suit of armour worn by his late father in order to defeat a vengeful spirit from the past.",
"A biopic detailing the 2 decades that Punjabi Sikh revolutionary Udham Singh spent planning the assassination of the man responsible for the Jallianwala Bagh massacre.",
"When a machine that allows therapists to enter their patients' dreams is stolen, all hell breaks loose. Only a young female therapist, Paprika, can stop it.",
"An ordinary word processor has the worst night of his life after he agrees to visit a girl in Soho whom he met that evening at a coffee shop.",
"A story that revolves around drug abuse in the affluent north Indian State of Punjab and how the youth there have succumbed to it en-masse resulting in a socio-economic decline.",
"A world-weary political journalist picks up the story of a woman's search for her son, who was taken away from her decades ago after she became pregnant and was forced to live in a convent.",
"Concurrent theatrical ending of the TV series Neon Genesis Evangelion (1995).",
"During World War II, a rebellious U.S. Army Major is assigned a dozen convicted murderers to train and lead them into a mass assassination mission of German officers.",
"The toys are mistakenly delivered to a day-care center instead of the attic right before Andy leaves for college, and it's up to Woody to convince the other toys that they weren't abandoned and to return home.",
"A soldier fighting aliens gets to relive the same day over and over again, the day restarting every time he dies.",
"After two male musicians witness a mob hit, they flee the state in an all-female band disguised as women, but further complications set in.",
"Exiled into the dangerous forest by her wicked stepmother, a princess is rescued by seven dwarf miners who make her part of their household.",
"A renegade reporter trailing a young runaway heiress for a big story joins her on a bus heading from Florida to New York, and they end up stuck with each other when the bus leaves them behind at one of the stops.",
"Story of 40-man Turkish task force who must defend a relay station.",
"Spinal Tap, one of England's loudest bands, is chronicled by film director Marty DiBergi on what proves to be a fateful tour.",
"Oskar, an overlooked and bullied boy, finds love and revenge through Eli, a beautiful but peculiar girl."]
descriptions_embeddings = list(
dense_embedding_model.embed(descriptions)
)
임베딩을 Qdrant에 업로드해 볼게요.
Qdrant 클라이언트는 간단한 인메모리 모드를 제공해서, 적은 양의 데이터로 로컬에서 실험할 수 있어요. 또는 실험용으로 Qdrant Cloud의 무료 클러스터를 이용할 수도 있어요.
from qdrant_client import QdrantClient, models
client = QdrantClient(":memory:") # Qdrant is running from RAM.
영화 데이터로 컬렉션을 만들어 볼게요.
client.create_collection(
collection_name="movies",
vectors_config={
"embedding": models.VectorParams(
size=client.get_embedding_size("sentence-transformers/all-MiniLM-L6-v2"),
distance=models.Distance.COSINE
)
}
)
그리고 임베딩을 업로드해요.
client.upload_points(
collection_name="movies",
points=[
models.PointStruct(
id=idx,
payload={"description": description},
vector={"embedding": vector}
)
for idx, (description, vector) in enumerate(
zip(descriptions, descriptions_embeddings)
)
],
)
임베딩을 암묵적으로 계산해서 업로드하기:
client.upload_points(
collection_name="movies",
points=[
models.PointStruct(
id=idx,
payload={"description": description},
vector={"embedding": models.Document(text=description, model=encoder_name)},
)
for idx, description in enumerate(descriptions)
],
)
1단계 검색 (First-stage retrieval)
all-MiniLM-L6-v2 기반 dense 검색기만 사용했을 때 결과가 얼마나 관련성 있는지 확인해 볼게요.
query = "A story about a strong historically significant female figure."
query_embedded = list(dense_embedding_model.query_embed(query))[0]
initial_retrieval = client.query_points(
collection_name="movies",
using="embedding",
query=query_embedded,
with_payload=True,
limit=10
)
description_hits = []
for i, hit in enumerate(initial_retrieval.points):
print(f'Result number {i+1} is \"{hit.payload["description"]}\"')
description_hits.append(hit.payload["description"])
임베딩을 암묵적으로 계산해서 쿼리하기:
query = "A story about a strong historically significant female figure."
initial_retrieval = client.query_points(
collection_name="movies",
using="embedding",
query=models.Document(text=query, model=encoder_name),
with_payload=True,
limit=10
)
결과는 다음과 같아요:
Result number 1 is "A world-weary political journalist picks up the story of a woman's search for her son, who was taken away from her decades ago after she became pregnant and was forced to live in a convent."
Result number 2 is "Exiled into the dangerous forest by her wicked stepmother, a princess is rescued by seven dwarf miners who make her part of their household."
...
Result number 9 is "A biopic detailing the 2 decades that Punjabi Sikh revolutionary Udham Singh spent planning the assassination of the man responsible for the Jallianwala Bagh massacre."
Result number 10 is "In 1431, Jeanne d'Arc is placed on trial on charges of heresy. The ecclesiastical jurists attempt to force Jeanne to recant her claims of holy visions."
가장 잘 맞는 *“The Messenger: The Story of Joan of Arc”*의 설명이 결과 10번째에 나오는 걸 볼 수 있어요.
이제 검색된 부분집합의 순서를 Jina Reranker v2로 정제해 볼게요. 리랭커는 쿼리와 문서(영화 설명) 집합을 입력받아, 쿼리와 각 문서 사이의 토큰 수준 상호작용을 바탕으로 관련성 점수를 계산해요.
new_scores = list(
reranker.rerank(query, description_hits)
) # returns scores between query and each document
ranking = [
(i, score) for i, score in enumerate(new_scores)
] # saving document indices
ranking.sort(
key=lambda x: x[1], reverse=True
) # sorting them in order of relevance defined by reranker
for i, rank in enumerate(ranking):
print(f'''Reranked result number {i+1} is \"{description_hits[rank[0]]}\"''')
리랭커는 원하던 영화를 관련성에 따라 첫 번째 위치로 옮겨줘요.
Reranked result number 1 is "In 1431, Jeanne d'Arc is placed on trial on charges of heresy. The ecclesiastical jurists attempt to force Jeanne to recant her claims of holy visions."
Reranked result number 2 is "Exiled into the dangerous forest by her wicked stepmother, a princess is rescued by seven dwarf miners who make her part of their household."
...
Reranked result number 9 is "An ordinary word processor has the worst night of his life after he agrees to visit a girl in Soho whom he met that evening at a coffee shop."
Reranked result number 10 is "A biopic detailing the 2 decades that Punjabi Sikh revolutionary Udham Singh spent planning the assassination of the man responsible for the Jallianwala Bagh massacre."
결론 (Conclusion)
리랭커는 검색된 후보를 더 깊은 의미 분석으로 재정렬해서 검색 결과를 정제해요. 효율을 위해서라면 검색된 부분집합에만 적용하는 게 핵심이에요.
리랭커의 힘을 활용해 검색의 속도와 정확성 사이의 균형을 잡아 보세요!