MTEB: Massive Text Embedding Benchmark에 관하여 (feat. MTEB-ko-retrieval Leaderboard, data contribution)

하루가 멀다하고 매일 새로운 모델들이 나오는 요즈음, 사용할 임베딩 모델을 고르는 것은 굉장히 피곤하다. 그래서 Open-LLM-leaderboard처럼 임베딩 모델들을 평가하여 공개하면 좋을 것 같다는 생각을 꾸준히 해왔다.

MTEB는 텍스트 임베딩 모델을 평가하기 위해 사용되는 라이브러리로, 다양한 태스크에 대해 다양한 데이터셋을 활용하여 굉장히 편리하게 평가할 수 있다. (딸깍) 이러한 측면에서 나의 니즈와 굉장히 잘 맞았고, MTEB를 유심히 살펴 보았다.

MTEB는 오픈 프로젝트이며, 누구나 새로운 데이터셋을 등록할 수 있고, 자유롭게 issue를 남겨 error fix 등을 할 수 있다. 오늘은 이와 관련된 나의 경험, 그리고 MTEB를 활용하여 한국어 검색 태스크를 수행한 코드 및 결과를 공유하고자 한다.

MTEB-ko-retrieval Leaderboard는 다음 Github repository에 공개되어 있다.

• issue에 평가하고 싶은 모델 남겨주시면 평가를 진행해 드리니, 필요하시면 남겨주세요 !! (⭐ 스타도 많이 눌러주세요 ⭐️)

[embed]GitHub - nlpai-lab/KURE: KURE: 고려대학교에서 개발한, 한국어 검색에 특화된 임베딩 모델 KURE: 고려대학교에서 개발한, 한국어 검색에 특화된 임베딩 모델. Contribute to nlpai-lab/KURE development by creating an account on GitHub.github.com

1. Evaluation

먼저 MTEB를 활용하여 embedding model들을 평가하는 코드부터 살펴보자.

1–1. Evaluation code

사실 평가는 굉장히 간단하다. 다음과 같이 MTEB 라이브러리를 활용한 코드를 실행하면 자동으로 평가가 진행된다. (코드는 Github에도 공개돼 있다.)

"""Benchmarking all datasets constituting the MTEB Korean leaderboard & average scores"""
from __future__ import annotations

import os
import logging
import argparse

import torch

from sentence_transformers import SentenceTransformer
from sentence_transformers.models import StaticEmbedding

import mteb
from mteb import MTEB, get_tasks
from dotenv import load_dotenv

from mteb.encoder_interface import PromptType
from mteb.models.sentence_transformer_wrapper import SentenceTransformerWrapper
from mteb.models.instruct_wrapper import instruct_wrapper

load_dotenv() # for OPENAI

parser = argparse.ArgumentParser(description="Extract contexts")
parser.add_argument('--quantize', default=False, type=bool, help='quantize embeddings')
args = parser.parse_args()

logging.basicConfig(level=logging.INFO)

logger = logging.getLogger("main")

TASK_LIST_CLASSIFICATION = []

TASK_LIST_CLUSTERING = []

TASK_LIST_PAIR_CLASSIFICATION = []

TASK_LIST_RERANKING = []

TASK_LIST_RETRIEVAL = [
    "Ko-StrategyQA",
    "AutoRAGRetrieval",
    "MIRACLRetrieval", # 시간이 오래 걸림 주의
    "PublicHealthQA",
    "BelebeleRetrieval",
    "MrTidyRetrieval", # 시간이 오래 걸림 주의
    "MultiLongDocRetrieval",
    "XPQARetrieval"
]

TASK_LIST_STS = []

TASK_LIST = (
    TASK_LIST_CLASSIFICATION
    + TASK_LIST_CLUSTERING
    + TASK_LIST_PAIR_CLASSIFICATION
    + TASK_LIST_RERANKING
    + TASK_LIST_RETRIEVAL
    + TASK_LIST_STS
)

model_names = [
    # my_model_directory
]
model_names = [
    "Salesforce/SFR-Embedding-2_R", # 4096
    "Alibaba-NLP/gte-Qwen2-7B-instruct", # 8192
    "BAAI/bge-multilingual-gemma2", # 8192
    "intfloat/e5-mistral-7b-instruct", # 32768
    "intfloat/multilingual-e5-large-instruct", # 512
    "openai/text-embedding-3-large", # 8191
    "Alibaba-NLP/gte-multilingual-base", 
    "intfloat/multilingual-e5-base", # 512
    "intfloat/multilingual-e5-large", # 512
    "jinaai/jina-embeddings-v3", # 8192
    "jhgan/ko-sroberta-multitask", # 128
    "BAAI/bge-m3", # 8192
    "nlpai-lab/KoE5", # 512
    "dragonkue/BGE-m3-ko", # 8192
    "Snowflake/snowflake-arctic-embed-l-v2.0" # 8192,
    "nlpai-lab/KURE-v1" # 8192
] + model_names

def evaluate_model(model_name):
    try:
        model = None
        if not os.path.exists(model_name): # hf에 등록된 모델의 경우
            if "m2v" in model_name: # model2vec 모델의 경우: 모델명에 m2v를 포함시켜주어야 model2vec 모델로 인식합니다.
                static_embedding = StaticEmbedding.from_model2vec(model_name)
                model = SentenceTransformer(modules=[static_embedding])
            else:
                if model_name == "nlpai-lab/KoE5":
                    # mE5 기반의 모델이므로, 해당 프롬프트를 추가시킵니다.
                    model_prompts = {
                        PromptType.query.value: "query: ",
                        PromptType.passage.value: "passage: ",
                    }
                    model = SentenceTransformerWrapper(model=model_name, model_prompts=model_prompts)
                elif model_name == "BAAI/bge-multilingual-gemma2":
                    # mbge-gemma2의 경우, mteb에서 지원하지 않습니다. 따라서, instruct_wrapper를 사용합니다.
                    instruction_template = '<instruct>{instruction}\n<query>'
                    model = instruct_wrapper(
                            model_name_or_path=model_name,
                            instruction_template=instruction_template,
                            attn="cccc",
                            pooling_method="lasttoken",
                            mode="embedding",
                            torch_dtype=torch.float16,
                            normalized=True,
                    )
                elif model_name == "Snowflake/snowflake-arctic-embed-l-v2.0":
                    # mteb에서 Snowflake 모델을 지원하지 않으므로, Snowflake에서 사용하는 "query: " prefix를 임의로 추가합니다.
                    model_prompts = {
                        PromptType.query.value: "query: ",
                    }
                    model = SentenceTransformerWrapper(model=model_name, model_prompts=model_prompts)
                else:
                    # mteb에 등록된 모델의 경우, 프롬프트/prefix 등을 포함하여 평가할 수 있습니다. 등록되지 않은 경우, sentence-transformer를 사용하여 불러옵니다.
                    model = mteb.get_model(model_name)
        else: # 직접 학습한 모델의 경우
            file_name = os.path.join(model_name, "model.safetensors")
            if os.path.exists(file_name):
                if "m2v" in model_name: # model2vec의 경우: 모델명에 m2v를 포함시켜주어야 model2vec 모델로 인식합니다.
                    static_embedding = StaticEmbedding.from_model2vec(model_name)
                    model = SentenceTransformer(modules=[static_embedding])
                else:
                    model = mteb.get_model(model_name)

        if model:
            print(f"Running task: {TASK_LIST} / {model_name}")
            evaluation = MTEB(
                tasks=get_tasks(tasks=TASK_LIST, languages=["kor-Kore", "kor-Hang", "kor_Hang"])
            )

            # 48GB VRAM 기준 적합한 batch sizes
            if "multilingual-e5" in model_name or "KoE5" in model_name:
                batch_size = 512
            elif "jina" in model_name:
                batch_size = 8
            elif "bge-m3" in model_name or "KURE" in model_name:
                batch_size = 32
            elif "gemma2" in model_name:
                batch_size = 256 
            elif "Salesforce" in model_name or "Snowflake" in model_name:
                batch_size = 128
            else:
                batch_size = 64

            evaluation.run(
                model,
                output_folder=f"results/{model_name}",
                encode_kwargs={"batch_size": batch_size},
            )
    except Exception as ex:
        print(ex)

if __name__ == "__main__":
    for model_name in model_names:
        evaluate_model(model_name)

• 만일 새로운 모델을 평가하고 싶다면 model_names 변수에 해당 모델의 이름을, 새로운 데이터셋으로 평가하고 싶다면 TASK_LIST 변수에 해당 데이터셋의 이름을 추가하면 된다. (Retrieval task가 아니어도 된다.) + 만일 위 코드에서 model2vec 변환을 거친 모델을 평가하고 싶다면 model 이름에 “m2v”를 넣어주면 된다.
• 또, 새로운 모델을 평가하는데 prompt/prefix를 넣어줘야 한다면, evaluate_model 함수에 작성한 것처럼 customize해서 사용할 수 있다.
• languages 변수에는 어떤 언어로 평가할지 설정할 수 있다. (MTEB의 각 task의 metadata-language를 보면 어떤 언어를 지원하는지 나와 있다.)
• 또, batch size를 높여 평가 시간을 단축할 수 있다. (모델마다 다름 주의)

1–2. Why MTEB?

그렇다면 왜 굳이 MTEB를 활용할까 ? 아무래도 가장 큰 장점은 prompt/ prefix의 반영과 편리함인 것 같다.

decoder-based retriever의 경우 prompt가 추가되어 학습되었으며, encoder-based retriever의 경우 prefix가 붙어 학습된 모델들이 있다. 예를 들어, Microsoft의 multilingual-e5 시리즈는 각 query 앞에 “query: ”, document 앞에 “passage: ” 라는 prefix가 붙어서 학습되었다.

이러한 모델들 중, MTEB에서 지원하는 오픈 모델들의 경우 (대부분 대기업의 multilingual model 들), 모델을 불러서 평가할 때 자동으로 prompt, prefix를 달아준다. 그러나 일부 모델 혹은 fine-tuning된 모델들은 지원하지 않는데, 그마저도 쉽게 달 수 있게 만들어두었다. (코드의 evaluate_model 함수 참고)

또, 평가하고자 하는 task, languages도 이름만 바꿔주면 평가할 수 있으며, 모든 지표에 대한 결과를 json 파일로 만들어주기 때문에, 굉장히 편리하다.

2. Dataset Contribution

MTEB는 dataset contribution 또한 굉장히 오픈되어 있다. 6월 정도까지는 데이터셋을 일정 이상 기여하고, MTEB repository에 많이 contribute 하면 MMTEB paper의 공동 저자로 넣어줬었는데, paper가 제출되어서 그런지 지금은 해당 활동을 진행하지 않는 것 같다 ㅠㅠ. 아무튼, 이번에는 데이터셋을 contribute 해본 경험을 작성해보도록 하겠다.

최근에 가장 잘 사용하고 있었던 임베딩 모델 평가 데이터셋 - Markers 회사의 AutoRAG에서 공개한 AutoRAG-example-korean-embedding-benchmark 를 MTEB repository에 올려보고자 하였다.

이 데이터는 원래 Allganize에서 한국어 RAG evaluation을 평가하기 위해 공개했던 데이터셋인데, AutoRAG 측에서 정제를 거쳐서 embedding model benchmark로 활용하였다. 공공기업의 문서를 페이지별로 PDF parsing 했다는 점에서 real-world scenario를 굉장히 잘 반영한 데이터라고 생각했다. (다만, retrieval pool이 작다는 단점이 있긴 하다.)

2–1. BEIR format

MTEB에 dataset을 추가하기 위해서는, 먼저 데이터셋을 BEIR format 으로 바꾸어야 한다. 즉 default, corpus, query의 3가지 subset이 있어야 한다. (데이터셋 예시를 보면 더 이해가 잘된다.)

• default에는 검색을 위한 query-id, corpus-id, score (score는 다 1)가,

• corpus에는 retrieval pool (corpus)에 해당하는 corpus-id와 text가,

• query에는 query에 해당하는 query-id와 text들이 있어야 한다.

2–2. Making an issue

이렇게 변환을 거친 후, MTEB Github repository에 issue를 남기면 된다. 내가 남긴 issue는 다음과 같다.

[embed]fix: Add Korean AutoRAGRetrieval by yjoonjang · Pull Request #1388 · embeddings-benchmark/mteb Adding datasets checklist Reason for dataset addition: This is a Korean Retrieval benchmark dataset covering 5 domains…github.com

가장 귀찮지만, MTEB에서 가장 중요하게 체크하는 것이 TaskMetadata이다. 이 항목은 해당 task의 metadata를 작성하는 란인데, 나는 다음과 같이 작성했다. Github code

# mteb/tasks/Retrieval/kor/AutoRAGRetrieval.py
from __future__ import annotations

from mteb.abstasks.TaskMetadata import TaskMetadata

from ....abstasks.AbsTaskRetrieval import AbsTaskRetrieval

class AutoRAGRetrieval(AbsTaskRetrieval):
    metadata = TaskMetadata(
        name="AutoRAGRetrieval",
        description="This dataset enables the evaluation of Korean RAG performance across various domains—finance, public sector, healthcare, legal, and commerce—by providing publicly accessible documents, questions, and answers.",
        reference="https://arxiv.org/abs/2410.20878",
        dataset={
            "path": "yjoonjang/markers_bm",
            "revision": "fd7df84ac089bbec763b1c6bb1b56e985df5cc5c",
        },
        type="Retrieval",
        prompt="Retrieve text based on user query.",
        category="s2p",
        modalities=["text"],
        eval_splits=["test"],
        eval_langs=["kor-Hang"],
        main_score="ndcg_at_10",
        date=("2024-08-03", "2024-08-03"),
        domains=["Government", "Medical", "Legal", "Social"],
        task_subtypes=["Article retrieval"],
        license="mit",
        annotations_creators="human-annotated",
        dialect=[],
        sample_creation="created",
        bibtex_citation="""@misc{kim2024autoragautomatedframeworkoptimization,
      title={AutoRAG: Automated Framework for optimization of Retrieval Augmented Generation Pipeline}, 
      author={Dongkyu Kim and Byoungwook Kim and Donggeon Han and Matouš Eibich},
      year={2024},
      eprint={2410.20878},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2410.20878}, 
}""",
    )

이렇게 작성하고 나면, 300M 크기의 작은 모델들 (intfloat/multilingual-e5-small, sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2)을 가지고 평가를 돌려보도록 시킨다. (TaskMetadata를 제대로 작성하지 않으면 평가도 안돌아간다 ㅠ)

이후, 결과가 괜찮게 나왔다면 리포트한 후, 리뷰를 요청하여 승인되면 dataset contribute가 마무리된다 !!

3. Error Resolving

마지막으로, MTEB로 다양한 모델들을 평가하면서 Error를 다룬 경험이 인상적이었어서, 이와 관련된 경험을 남기고 마무리한다.

먼저 내가 마주한 에러는 openai의 text-embedding-3-large 모델을 평가할 때, 임베딩하려는 text가 openai-text-embedding이 지원하는 최대 길이인 8191을 넘어서 발생하는 문제였다. 대부분의 오픈 모델들은 알아서 잘라 주는데, api로 제공되어서 그런지 알아서 잘라주지 못하고 다음 에러를 반환했다:

ERROR:mteb.evaluation.MTEB:Error while evaluating MIRACLRetrieval: Error code: 400 - {'error': {'message': "This model's maximum context length is 8192 tokens, however you requested 8461 tokens (8461 in your prompt; 0 for the completion). Please reduce your prompt; or completion length.", 'type': 'invalid_request_error', 'param': None, 'code': None}}
Error code: 400 - {'error': {'message': "This model's maximum context length is 8192 tokens, however you requested 8461 tokens (8461 in your prompt; 0 for the completion). Please reduce your prompt; or completion length.", 'type': 'invalid_request_error', 'param': None, 'code': None}}

곧바로 이를 issue로 남겼고, 작업을 시작했다.

[embed]fix: add sentence trimming to OpenAIWrapper by yjoonjang · Pull Request #1526 ·… 1494 @Samoedgithub.com

사실 뭐 큰 문제는 아니고, 그냥 text-embedding-3-large가 사용하는 tokenizer를 불러와 sentence를 trim 하고, 다시 api로 넘겨주면 되는 문제였다. 작업을 하고 리뷰를 요청했는데, 생각보다 굉장히 빠르고 열심히 리뷰해줘서 고마웠다.

위 issue를 열고 진행하면서, 같은 문제를 겪고 있던 MTEB maintainer (Muennighoff)의 코드를 일부 반영하여 수정하기도 하고, tiktoken 때문에 발생하는 에러로 인해 openai를 위한 라이브러리 버전 관리도 해보았다.

회사를 다니지 않으면 이런 경험은 쉽게 해보지 못할 것 같은데, 그러한 점에서 작지만 굉장히 인상적인 경험이었다.

이렇게 오늘은 MTEB에 대한 간단한 소개와 MTEB에 contribution을 남긴 나의 경험들에 대해 정리해 보았다. 어렵지 않은 작업이라, 관심 있는 분들은 한번 해보신다면 좋을 것 같다.

아 그리고 MTEB를 활용한 MTEB-ko-retrieval Leaderboard도 출시했으니 많은 관심(⭐ 스타 ⭐) 부탁드립니다 !!