embedding-train

embedding-train 프로젝트는 multilingual-e5 모델을 한국어 데이터로 fine-tuning하여 한국어에 특화된 임베딩 모델을 만드는 것을 목표로 합니다. 이를 통해 한국어 질의와 문서 간의 의미적 유사성을 잘 반영할 수 있는 고성능 임베딩 모델을 구축합니다.

프로젝트 버전

v1: InfoNCE loss를 huggingface Trainer의 compute_loss 함수에 직접 구현하여 학습하는 파이프라인입니다. 이를 통해 기본적인 한국어 임베딩 학습이 가능합니다. 관련 블로그
v1.1: InfoNCE loss가 배치 크기에 큰 영향을 받는다는 사실을 기반으로, gradient cache를 적용하여 큰 배치 크기에서도 안정적으로 학습할 수 있도록 개선된 파이프라인입니다.

학습

모델을 파인 튜닝하고 학습시키기 위한 단계는 다음과 같습니다:

의존성 패키지를 설치 합니다.

pip install -r requirements.txt

v1.1/scripts/finetune.sh 스크립트를 수정하여 학습 파라미터를 설정합니다.
설정한 파라미터에 따라 스크립트를 실행하여 모델을 학습시킵니다.

bash v1.1/scripts/finetune.sh

데이터 구성

모델 학습에 사용된 데이터는 다음과 같은 구조를 가지고 있습니다:

{
    "query": "사용자가 입력한 질의",
    "document": "질의와 관련된 문서",
    "hard_negative": "질의와 관련성이 적지만, 유사해 보이는 문서"
}

데이터는 AIHUB, KorQUAD, KommonGen, Exobrain, KLUE, KoBEST, NIKL로부터 총 10종의 오픈 데이터를 수집하였으며, 총 데이터 통계는 다음과 같습니다.

데이터 원천	# 샘플
AIHUB	552,276
KorQuAD	28,795
KommonGen	14,893
Exobrain	78,362
KLUE	6,334
KoBEST	2,849
NIKL	39,919
Total	723,428

평가

python evaluate.py명령어로 평가를 수행합니다.
streamlit run leaderboard.py명령어로 평가에 대한 리더보드를 확인합니다.

결과

Ko-strategyQA, AutoRAG-embedding-benchmark, PublicHealthQA의 총 3가지 평가 데이터셋으로 평가를 진행했으며, 결과는 다음과 같습니다.

Model	Ko-strategyQA				AutoRAG-benchmark				PublicHealthQA				Avg
Model	NDCG@1	F1@1	NDCG@3	F1@3	NDCG@1	F1@1	NDCG@3	F1@3	NDCG@1	F1@1	NDCG@3	F1@3	NDCG@1	F1@1	NDCG@3	F1@3
BGE-m3	75.68	59.68	74.79	55.76	67.54	67.54	47.57	43.42	67.53	67.53	76.69	41.56	70.25	64.92	66.35	46.91
mGTE	70.1	54.91	70.06	52.55	58.77	58.77	46.14	40.35	58.44	58.44	69.06	38.31	62.44	57.37	61.75	43.74
mE5-large	76.86	61.07	76.48	57.05	63.15	63.15	44.04	39.91	68.83	68.83	79.31	42.86	69.61	64.35	66.61	46.61
kf-deberta	60.64	47.19	61.21	46.22	45.61	45.61	36.79	31.58	54.55	54.55	64.69	35.71	53.6	49.12	54.23	37.84
KoE5 (Ours)	76.69	60.70	75.70	56.32	70.17	70.17	48.01	44.30	71.43	71.43	80.10	42.86	72.76	67.43	67.94	47.83

yjoonjang/embedding-train

embedding-train

프로젝트 버전

학습

데이터 구성

평가

결과