[LLM] RAG 개선하기 (2) - 임베딩 모델 미세 조정하기 (Fine-Tuning)

 
이번 시간에는 좋은 임베딩 모델을 만들어 검색 쿼리와 관련된 문서를 잘 찾게 만들어 보자. 이를 위해 임베딩 모델을 KLUE의 MRC 데이터셋으로 추가 학습시키도록 하겠다.

---

사전 학습된 모델의 성능 살펴보기

 

데이터 전처리

 

from datasets import load_dataset
klue_mrc_train = load_dataset('klue', 'mrc', split='train')
klue_mrc_train[0]
# {'title': '제주도 장마 시작 … 중부는 이달 말부터',
#  'context': '올여름 장마가 17일 제주도에서 시작됐다. 서울 등 중부지방은 예년보다 사나흘 정도 늦은 이달 말께 장마가 시작될 전망이다.17일 기상청에 따르면 제주도 남쪽 먼바다에 있는 장마전선의 영향으로 이날 제주도 산간 및 내륙지역에 호우주의보가 내려지면서 곳곳에 100㎜에 육박하는 많은 비가 내렸다. 제주의 장마는 평년보다 2~3일, 지난해보다는 하루 일찍 시작됐다. 장마는 고온다습한 북태평양 기단과 한랭 습윤한 오호츠크해 기단이 만나 형성되는 장마전선에서 내리는 비를 뜻한다.장마전선은 18일 제주도 먼 남쪽 해상으로 내려갔다가 20일께 다시 북상해 전남 남해안까지 영향을 줄 것으로 보인다. 이에 따라 20~21일 남부지방에도 예년보다 사흘 정도 장마가 일찍 찾아올 전망이다. 그러나 장마전선을 밀어올리는 북태평양 고기압 세력이 약해 서울 등 중부지방은 평년보다 사나흘가량 늦은 이달 말부터 장마가 시작될 것이라는 게 기상청의 설명이다. 장마전선은 이후 한 달가량 한반도 중남부를 오르내리며 곳곳에 비를 뿌릴 전망이다. 최근 30년간 평균치에 따르면 중부지방의 장마 시작일은 6월24~25일이었으며 장마기간은 32일, 강수일수는 17.2일이었다.기상청은 올해 장마기간의 평균 강수량이 350~400㎜로 평년과 비슷하거나 적을 것으로 내다봤다. 브라질 월드컵 한국과 러시아의 경기가 열리는 18일 오전 서울은 대체로 구름이 많이 끼지만 비는 오지 않을 것으로 예상돼 거리 응원에는 지장이 없을 전망이다.',
#  'news_category': '종합',
#  'source': 'hankyung',
#  'guid': 'klue-mrc-v1_train_12759',
#  'is_impossible': False,
#  'question_type': 1,
#  'question': '북태평양 기단과 오호츠크해 기단이 만나 국내에 머무르는 기간은?',
#  'answers': {'answer_start': [478, 478], 'text': ['한 달가량', '한 달']}}

 
먼저 KLUE의 MRC 데이터셋을 받아오고 구조를 살펴보자. 데이터 구조에서 question과 context는 질문 쿼리와 응답에 해당한다. 이 데이터셋을 이용하여 사전 학습된 모델을 미세 조정한다면, 이 데이터에 해당하는 질문에 대한 유사도 높은 답변을 제공할 수 있을 것이다.
 

from datasets import load_dataset
klue_mrc_train = load_dataset('klue', 'mrc', split='train')
klue_mrc_test = load_dataset('klue', 'mrc', split='validation')

df_train = klue_mrc_train.to_pandas()
df_test = klue_mrc_test.to_pandas()

df_train = df_train[['title', 'question', 'context']]
df_test = df_test[['title', 'question', 'context']]

 
데이터셋을 훈련, 테스트 데이터셋으로 분리해 주고, 데이터의 전처리를 쉽게 하기 위해 pandas dataframe으로 변환한다. 이어서 우리에게 필요한 세 가지 column만 남겨두고 나머지는 제거해 주었다.
 

def add_ir_context(df):
  irrelevant_contexts = []
  for idx, row in df.iterrows():
    title = row['title']
    irrelevant_contexts.append(df.query(f"title != '{title}'").sample(n=1)['context'].values[0])
  df['irrelevant_context'] = irrelevant_contexts
  return df

df_train_ir = add_ir_context(df_train)
df_test_ir = add_ir_context(df_test)

 
이어서 관련 없는 데이터셋을 만들어 보자. 이걸 하는 이유는 현재 KLUE MRC 데이터셋에는 서로 대응되는 질문-내용 쌍만 존재하기 때문에 학습한 모델이 문장의 유사도를 잘 계산하는지 평가할 수가 없다. 따라서 일부러 그런 데이터셋을 만들어 주는 것이다. 각 행을 돌면서, 해당 행 데이터의 title과 다른 title을 가진 데이터의 context를 무작위로 골라 데이터프레임의 irrelevant_context에 추가한다. 훈련 데이터셋으로 하나, 테스트 데이터셋으로 하나를 만들어 준다.
 

from sentence_transformers import InputExample

examples = []
for idx, row in df_test_ir[:100].iterrows():
  examples.append(
      InputExample(texts=[row['question'], row['context']], label=1)
  )
  examples.append(
      InputExample(texts=[row['question'], row['irrelevant_context']], label=0)
  )

 
이제 마지막 단계이다. 사전 학습된 모델에 미세 조정 없이 정확도를 확인하기 위해, 우리가 만든 데이터셋에 관련 있는 context가 들어있는 경우 1을, 관련 없는 context가 들어있는 경우 0을 라벨링 해준다.
 

사전 학습된 모델 성능 측정

 

from sentence_transformers import SentenceTransformer
sentence_model = SentenceTransformer('shangrilar/klue-roberta-base-klue-sts')

 
모델을 불러오고,
 

from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator

evaluator = EmbeddingSimilarityEvaluator.from_input_examples(examples)
evaluator(sentence_model)

# 0.8151553052035344

 
examples를 바탕으로 사전학습된 모델의 성능을 측정해 보면 0.815의 성능이 나오는 것을 확인할 수 있다.
 

MNR 손실을 활용해 미세 조정하기

 
이번에는 MNR 손실을 활용하여 미세 조정을 통해 임베딩 모델의 성능을 올려 보자. Multiple Negatives Ranking의 특징은 정답 데이터만을 학습에 사용하여도 하나의 배치 데이터 안에서 다른 데이터의 context를 관련이 없는 데이터로 사용해 모델을 학습시키므로 서로 관련이 있고 없고를 나누기에 좋은 학습 손실이다.
 

train_samples = []
for idx, row in df_train_ir.iterrows():
    train_samples.append(InputExample(
        texts=[row['question'], row['context']]
    ))

 
따라서 MNR 손실을 활용해 모델을 학습시킬 경우에는 정답 데이터만을 이용해 주면 된다.
 

from sentence_transformers import datasets

batch_size = 16

loader = datasets.NoDuplicatesDataLoader(
    train_samples, batch_size=batch_size)

 
중복데이터를 알아서 제거해 주는 NoDuplicateDataLoader를 이용해 한번 걸러주고, 남은 데이터르 16의 배치 사이즈로 담는다.
 

from sentence_transformers import losses

loss = losses.MultipleNegativesRankingLoss(sentence_model)

 
라이브러리를 통해 MNR 손실을 가져와 주고,
 

epochs = 1
save_path = './klue_mrc_mnr'

sentence_model.fit(
    train_objectives=[(loader, loss)],
    epochs=epochs,
    warmup_steps=100,
    output_path=save_path,
    show_progress_bar=True
)

 
.fit() 메서드를 통해 학습을 진행한다.
 

evaluator(sentence_model)
# 0.8600968992433692

 
미세 조정이 끝난 모델에 대해 평가를 해보면 기존의 0.815의 성능에 비해 0.860의 성능 향상을 확인할 수 있었다!