Ko-MTEB 따라잡기 (작성 중) - Embedding

학습 공부를 제대로 시작하기전에, 기존 모델들 부터 재현해보면서 Evaluation하는 것을 시작해보려고 한다.

(회사에서 이걸 과연.. 지켜봐줄지는 모르겠지만...)

최종 목표는 Embedding/Rerank 모델을 pre-train, fine-tune, evaluation까지, 그리고 HF에도 올려보고..

솔직히 성능 향상을 보진 못할 것 같다. 데이터 수집부터 정제, 학습까지 다 혼자서 진행해야할 것 같아서 이게 내가 잘 할 수 있는 것인진 모르겠다. 최선을 다해보자

사실 많은 사람들이 MTEB 벤치마크를 커스텀한 코드들을 많이 공유했다.

하지만 왜 직접하려고하냐면, 내 코드가 아니기 때문이다.. 내 코드가 아니니까 제대로 이해가 안되고, 평가를 진행하다보면 어떻게 학습을 해야하는지, 어떤 Output이 나오고를 알 수 있을거라 생각했다.

최종 목표는 ColBERT MultiVector Model Evaluation까지 생각하고 있다. (MUVERA 포함)

따라서, Ko-MTEB, 재구현한다!

MTEB

Massive Text Embedding Benchmark (MTEB)

•

텍스트 임베딩 모델의 평가 한계를 극복 → 8개 Task를 포괄하는 벤치마크 제시

•

latency / emb size와의 trade-off 분석까지도 가능

회사에서 가장 집중하고 있는 부분은 RAG부분이기에, 이 중 우선 Retrieval에 대한 평가에 대해 집중적으로 진행하려고 한다. (Retrieval은 nDCG@k를 기반으로 평가하려고한다)

Retrieval task를 IR (information Retrieval) 이라고도 한다

Benchmark Code

정말 감사하게도 많은 분들께서 K-MTEB Evaluation과 관련되어 정보를 공유해주셨고, 아래 내용들을 참조했다.

GitHub - BM-K/Korean-MTEB-Retrieval-Evaluators: Korean MTEB Retrieval Evaluators for SPLADE, Dense, and Reranking models

Korean MTEB Retrieval Evaluators for SPLADE, Dense, and Reranking models - BM-K/Korean-MTEB-Retrieval-Evaluators

github.com

https://medium.com/@yjoonjang/mteb-massive-text-embedding-benchmark에-관하여-feat-mteb-ko-retrieval-leaderboard-data-contribution-f4f9c979175d

medium.com

임베딩 모델 평가) MTEB 코드 살펴보기 (2) - Customizing

아래 설명처럼 평가된 한국어 임베딩 모델 순위 확인하기 : https://github.com/OnAnd0n/ko-embedding-leaderboard 이전 글 : 임베딩 모델 평가) MTEB 코드 살펴보기 (1) 다음 글: 임베딩 모델 평가) 한국어 임베딩 모델 리더보드 (3)앞선 글에서는 MTEB 깃허브의 Example 코드로 Model / Tasks / Evaluation이 어떻게 수행되는지,코드를 중심으로 확인하였다. 이전 글에서는 Model이 Sentence-Transformers로 호출될 때의 코드 흐름을 확인하였는데, MTEB에서는 Sentence-Transformers가 아닌 다른 방식으로도 평가를 수행할수있다.그리고 이에더하여 Customizing하여 모델을 평가할수 있도록 지원하기도..

introduce-ai.tistory.com

기본적으로 SBERT기반의 모델들은 바로 사용할 수 있다.

하지만 개별적으로 만든 모델이라면 Encoder Interface를 사용해서 상속받아 구현해야한다.

•

Query Fix가 있는 경우 PromptType 에 맞춰서 query, passage를 지정해주어야 작동한다고 한다.

근데 이건 어디서 찾아야하는지 모르겠다.. 보통 HuggingFace model Card에 있을 줄 알았는데..

Evaluation Model Select

그냥 깡으로는 코드짜기가 어려울 것 같아서, 우선은 몇개 모델을 선정해서 Evaluation을 수행해보자

모델선정은 Korean Specific으로 설정하고, 그 중 base 모델들이 다른 것들을 기준으로 선정했다. (only open)

테스트 해볼 모델들

•

ColBERT 계열

◦

yjoonjang/colbert-ko-v1.0

•

ModernBERT 계열

◦

SK A.X Encoder? → 주로 SPLADE로 말아서 사용?

•

e5 계열

◦

Query Prefix 존재

◦

intfloat//multilingual-e5-large-instruct

◦

PwC_Embedding_expr

•

BGE-m3 계열

◦

BAAI/BGE-m3

◦

dragonkue/BGE-m3-ko

•

XLM-RoBERTa 계열 (snowflake)

◦

Query Prefix 존재

◦

dragonkue/snowflake-arctic-embed-l-v2.0-ko

•

Decoder

◦

Qwen3-Embedding-0.6b

Evaluation Tasks Select

tasks = mteb.get_tasks(languages=["kor"]) 로 테스트 가능한 kor-tasks 들을 출력해보았다.

그 중 다음과 같이 분류할 수 있다. (Vision Task 제외)

Retrieval (문서/지식 검색)

Reranking (재정렬)

Bitext Mining (병렬문장 매칭)

STS (문장 의미 유사도)

Classification (문장/문서 분류)

NLI (Natural Language Inference / Entailment)

Classification (문장/문서 분류)

Clustering (군집화)

이 중 Retrieval로 테스트 진행하고 STS, Classification, NLI, Clustering 순으로 진행하려고 한다.

Simple Test

첫 모델은 dragonkue/snowflake-arctic-embed-l-v2.0-ko 로 선정하여 테스트하고, Retrieval Task는 6개로 테스트를 진행한다.

Task들은 KoStrategyQA, AutoRAGRetrieval, XPQARetrieval, BelebeleRetrieval, MultiLongDocRetrieval, PublicHealthQARetrieval 총 6개를 기반으로 테스트를 진행한다.

Retrieval에서 사용할 수 있는 kor-task들은 mteb.get_tasks(languages=["kor"], task_types=["Retrieval"]) 명령어로 확인할 수 있다.

from util import get_model_infos

TASK_RETRIEVAL = [
    # "Ko-StrategyQA",
    "AutoRAGRetrieval",
    # "XPQARetrieval",
    # "BelebeleRetrieval",
    # "MultiLongDocRetrieval",
    # "PublicHealthQA"
]

def test_evaluation(model_name, tasks):
    model = None
    model_card_info = get_model_infos(model_name)
    if model_card_info["eval_library"] == "sentence-transformers":
        # if model_name == "dragonkue/snowflake-arctic-embed-l-v2.0-ko":
        model = SentenceTransformer(model_name)
        bch = 3
    evaluation = MTEB(tasks=get_tasks(tasks=tasks, languages=["kor-Kore", "kor-Hang", "kor_Hang"]))
    evaluation.run(
        model=model,
        output_folder=f"results/{model_name}",
        encode_kwargs={"batch_size":bch}
    )

if __name__ == "__main__":
    model_name = "dragonkue/snowflake-arctic-embed-l-v2.0-ko"
    test_evaluation(model_name, TASK_RETRIEVAL)

nlpai-lab/KURE 코드를 참고하여 evaluation test를 수행해보았다.

KURE의 코드의 경우 MTEB evaluation을 자동으로 평가해주도록 코드가 작성되어있고,
BM-K님 코드의 경우 BEIR 기반으로 Evaluate 코드가 작성되어있다. (추후에 다시 보기..)
이번에는 KURE 코드 기반으로 시도하려고 한다.

ColBERT Model Evaluation

최근 ColBERT로 Multi-Embedding → MUVERA로

기존 KURE 코드에서는, ColBERT MultiEmbedding Model에 관해서는 평가하는 코드가 없다.

test model : yjoonjang/colbert-ko-v1.0

MTEB에서 Pylate기반으로 Wrapper를 제공해주는데, ColBERTWrapper 를 사용해서 평가하면된다.

class ColBERTWrapper(Wrapper):
    def __init__(
        self,
        model_name: str,
        revision: str | None = None,
        model_prompts: dict[str, str] | None = None,
        **kwargs,
    ) -> None:
        requires_package(self, "pylate", model_name, "pip install mteb[pylate]")
        from pylate import models as colbert_model  # type: ignore[import]

        self.model_name = model_name
        self.model = colbert_model.ColBERT(self.model_name, revision=revision, **kwargs)
        built_in_prompts = getattr(self.model, "prompts", None)
        if built_in_prompts and not model_prompts:
            model_prompts = built_in_prompts
        elif model_prompts and built_in_prompts:
            logger.info(f"Model.prompts will be overwritten with {model_prompts}")
            self.model.prompts = model_prompts
        self.model_prompts = self.validate_task_to_prompt_name(model_prompts)
  .....

ColBERT는 late interaction 기반 프레임워크인 pylate를 사용하고 colbert_model.ColBERT 에서 불러온 모델을 사용한다.

따라서 해당 wrapper를 사용하면, 아래와 같이 평가할 수 있다.

def ColBERT_test_Evaluation(model_name, tasks, device_number):
    
    model = ColBERTWrapper(
        model_name=model_name,
        model_kwargs={
            "torch_dtype": torch.float32 # f32
        },
        device=torch.device(f"cuda:{str(device_number)}")
    )
    
    if hasattr(model.model[0].tokenizer, "model_max_length"):
        model.model[0].tokenizer.model_input_nmaes = ["input_ids", "attention_mask"] # token_type_ids가 필요없음
        
    evaluation = MTEB(tasks=get_tasks(tasks=tasks, languages=["kor-Kore", "kor-Hang", "kor_Hang"]))
    evaluation.run(
        model=model,
        output_folder=f"results/{model_name}",
        encode_kwargs={"batch_size": 16}
    )

나머지 모델들은 sentence Transformer에서 호출해서 사용하면 될 것 같다.

get_model_info 함수를 만들어서 pylate가 필요한 경우에는 다르게 Wrapping해야하는가?