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 | ||||||||||||
---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
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 |