11. 자신의 데이터에 맞춘 임베딩 모델 만들기 : RAG 개선하기

 

 

☀️  Summary 

❍  바이인코더 + 교차인코더 결합을 통한 RAG 성능 향상 
❍  사전학습된 언어모델 + Pooling층 + 파인튜닝 

 

 

 

 

1.   검색 성능을 높이기 위한 두 가지 방법 

---

 

❍  바이인코더와 교차인코더를 결합해 더 빠른 검색을 수행 

 ↪︎  1) 바이인코더를 사용해 대규모의 문서에서 검색 쿼리와 유사한 소수의 문서(ex.상위 100개)를 선별 

 ↪︎  2) 의미검색을 통해 선별한 소수의 문서는 유사도를 더 정확히 계산할 수 있는 교차 인코더를 사용해 유사한 순서대로 재정렬 (Rerank) 

 

https://cohere.com/blog/rerank

 

 

❍  바이인코더와 교차인코더를 결합할 때, 검색 성능을 더 높이는 방법

 ↪︎  1) 바이인코더 추가 학습 : 문장 임베딩 모델도 학습 데이터와 유사한 입력에 대해 잘 작동하므로, 사용하려는 데이터셋으로 추가학습을 시킨다. 

 ↪︎  2) 교차인코더 추가 학습 :  검색된 모든 문서가 아니라, 상위 몇 개의 입력만 프롬프트에 추가해 검색 증강 생성이 효과적으로 작동하도록 함 

 

 

 

 

2.   언어모델을 임베딩 모델로 만들기 

---

 

 

↪︎  문장 임베딩 모델은 2개의 층으로 나뉜다. 첫 번째 층은 대량의 텍스트 데이터로 사전학습한 BERT나 RoBERTa 같은 언어모델이다. 두번째 층은 풀링층으로 입력 문장의 길이에 따라 달라질 수 있는 출력 차원을 고정된 차원으로 맞추는 역할을 한다. 풀링의 방식에는 클래스모드, 평균모드, 최대모드가 있는데 일반적으로 평균 모드를 많이 사용한다. 이처럼 사전학습된 언어모델을 불러오고 그 위에 풀링 층을 추가하고 문장의 의미를 잘 담을 수 있도록 학습해야 한다. 

↪︎  Sentence-Tranformers 라이브러리를 사용하면 문장 임베딩 모델을 쉽게 활용할 수 있다. 

 

 

2.1  대조학습 

 

❍  Contrastive learning

 ↪︎  관련이 있거나 유사한 데이터는 더 가까워지도록 만들고 관련이 없거나 유사하지 않은 데이터는 더 멀어지도록 하는 학습 방식

 ↪︎  대조학습 방식으로 임베딩 모델을 학습시킬 때, 다양한 데이터를 사용할 수 있다. 2개의 문장을 임베딩 모델에 각각 입력하고 서로 유사한 데이터인 경우는 가깝게 그렇지 않은 경우는 멀게 만들 수 있다. 또는 서로 이어지는 문장이라면 가깝게 그렇지 않으면 서로 멀게 만들 수 있다. 마지막으로 서로 질문답변 관계인 경우라면 가깝게 아니면 멀게 학습시킬 수 있다. 

 

 

 

2.2  실습 [1] 

 

❍  언어모델을 그대로 불러와 문장 임베딩을 만들어보기 

 

1) 사전 학습된 언어모델을 불러와 문장 임베딩 모델 만들기 

 

from sentence_transformers import SentenceTransformer, models

# 1) modules 모듈을 활용해 사전학습된 모델 불러오기 
transformer_model = models.Transformer('klue/roberta-base')

# 2) 평균 풀링층 생성 
pooling_layer = models.Pooling(
    transformer_model.get_word_embedding_dimension(),
    pooling_mode_mean_tokens=True
)

# 3) SentenceTransformer 클래스로 두 모듈을 결합해 문장 임베딩 모델 생성 
embedding_model = SentenceTransformer(modules=[transformer_model, pooling_layer])

 

 

 

 

2) 실습 데이터셋 불러오기

 ↪︎  KLUE의 STS 데이터셋 (두 문장이 서로 얼마나 유사한지 점수를 매긴 데이터셋) 

 ↪︎  labels 칼럼에 두 문장이 얼마나 유사한지를 나타내는 다양한 형식의 레이블이 있음 (label을 사용할 예정) 

 

from datasets import load_dataset
klue_sts_train = load_dataset('klue', 'sts', split='train')
klue_sts_test = load_dataset('klue', 'sts', split='validation')
klue_sts_train[0]

# {'guid': 'klue-sts-v1_train_00000',
#  'source': 'airbnb-rtt',
#  'sentence1': '숙소 위치는 찾기 쉽고 일반적인 한국의 반지하 숙소입니다.',
#  'sentence2': '숙박시설의 위치는 쉽게 찾을 수 있고 한국의 대표적인 반지하 숙박시설입니다.',
#  'labels': {'label': 3.7, 'real-label': 3.714285714285714, 'binary-label': 1}}




# 학습 데이터셋의 10%를 검증 데이터셋으로 구성 
klue_sts_train = klue_sts_train.train_test_split(test_size=0.1, seed=42)
klue_sts_train, klue_sts_eval = klue_sts_train['train'], klue_sts_train['test']

 

 

 

3) label 정규화하기 

 

from sentence_transformers import InputExample


# 유사도 점수를 0~1 사이로 정규화 하고 InputExample 객체에 담기 
def prepare_sts_examples(dataset):
    examples = []
    for data in dataset:
        examples.append(
            InputExample(
                texts=[data['sentence1'], data['sentence2']], # 텍스트쌍을 리스트 형태로 입력
                label=data['labels']['label'] / 5.0) # 정규화 
            )
    return examples


# 데이터 전처리 수행 
train_examples = prepare_sts_examples(klue_sts_train)
eval_examples = prepare_sts_examples(klue_sts_eval)
test_examples = prepare_sts_examples(klue_sts_test)

 

 

4) 배치 데이터셋 만들기 

 

from torch.utils.data import DataLoader
train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=16)

 

 

 

5) 검증을 위한 평가 객체 준비 및 모델 성능 결과 

 ↪︎  EmbeddingSimilarityEvaluator 를 사용해 임베딩 모델의 성능을 평가할 때 사용할 수 있도록 준비 

 ↪︎  from_input_examples 메서드를 사용해, 검증 데이터셋과 평가 데이터셋을 사용하는 평가 객체를 생성 

 

from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator

eval_evaluator = EmbeddingSimilarityEvaluator.from_input_examples(eval_examples)
test_evaluator = EmbeddingSimilarityEvaluator.from_input_examples(test_examples)


test_evaluator(embedding_model)
# 0.36460670798564826

 

 ↪︎  언어모델을 그대로 문장 임베딩 모델로 만든 embedding_model이 얼마나 문장의 의미를 잘 반영해 문장 임베딩을 생성하는지 확인해보았을 때, 0.364로 역할을 잘하지 못하는 것을 확인할 수 있다. 

 

 

 

 

2.3  실습 [2] 유사한 문장 데이터로 임베딩 모델 학습하기 

 

❍  임베딩 모델 학습

 ↪︎  기본 언어모델 :  klue/roberta-base

 ↪︎  학습에 사용될 손실함수 : CosineSimilarityLoss를 사용한다. 학습 데이터를 문장 임베딩으로 변환하고 두 문장 사이의 코사인 유사도와 정답 유사도를 비교해 학습을 수행한다. 

 

from sentence_transformers import losses

num_epochs = 4
model_name = 'klue/roberta-base'
model_save_path = 'output/training_sts_' + model_name.replace("/", "-")
train_loss = losses.CosineSimilarityLoss(model=embedding_model)

# 임베딩 모델 학습
embedding_model.fit(
    train_objectives=[(train_dataloader, train_loss)],
    evaluator=eval_evaluator,
    epochs=num_epochs,
    evaluation_steps=1000,
    warmup_steps=100,
    output_path=model_save_path
)

 

 

 

❍  학습한 임베딩 모델의 성능평가 

 ↪︎  학습된 모델이 저장된 경로에서 모델을 읽어오고, 임베딩 모델을 평가한다. 점수가 0.364에서 0.896으로 크게 향상된 것을 볼 수 있다. 

trained_embedding_model = SentenceTransformer(model_save_path)
test_evaluator(trained_embedding_model)
# 0.8965595666246748

 

 

 

❍  허깅페이스 허브에 모델 저장 

 ↪︎  계정 토큰을 통해 허깅페이스 허브에 접근해, 모델을 업로드한다. 

from huggingface_hub import login
from huggingface_hub import HfApi

login(token='허깅페이스 허브 토큰 입력')
api = HfApi()
repo_id="klue-roberta-base-klue-sts"
api.create_repo(repo_id=repo_id)


# 모델 업로드 
api.upload_folder(
    folder_path=model_save_path,
    repo_id=f"본인의 허깅페이스 아이디 입력/{repo_id}",
    repo_type="model",
)

 

 

 

 

3.   임베딩 모델 미세 조정하기 

---

 

3.1  학습 준비

 

❍ [복습] RAG

 ↪︎  RAG는 검색 쿼리와 관련된 문서를 찾아, LLM 프롬프트에 맥락 데이터로 추가할 때, 임베딩 모델을 활용한다. 좋은 임베딩 모델이라면 검색 쿼리와 관련있는 문서는 유사도가 1에 가깝게 나와야 한다. 

 ↪︎  임베딩모델을 KLUE의 MRC 데이터셋(기사 본문 및 해당 기사와 관련된 질문을 수집한 데이터)으로 추가학습시켜 실습 데이터의 문장 사이의 유사도를 더 잘 계산할 수 있도록 만든다. 

 

 

❍  미세조정 

 ↪︎  문장 임베딩 모델도, 다른 딥러닝 모델과 마찬가지로 학습 데이터와 유사한 데이터일 때 가장 잘 동작한다. 

 ↪︎  따라서, 사전학습된 임베딩 모델을 그대로 활용하는 경우, 사전 학습에 사용된 데이터셋이, 실습에 사용하는 MRC 데이터셋과 단어, 주제 등이 다르면 성능이 낮아진다. 따라서 MRC 데이터셋에 임베딩 모델을 활용하려고 한다면, 그 목적에 맞게 MRC 데이터셋으로 미세조정 해야 한다. 

 

 

❍  실습  

 

1) 2장에서 저장한 기본 임베딩 모델 불러오기 

 

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

 

 

2) 질문과 관련없는 기사 추가하기 

 ↪︎  question - context 칼럼은 관련이 있으므로 label 1을 부여하고, 임의로 질문과 관련없는 기사를 추가해 만든 irrelevant_context 칼럼은 관련이 없으므로 label 0을 부여

 

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)

 

 

 

3) 기본 임베딩 모델을 미세조정하지 않은 상태에서 성능 평가 

 

from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator
evaluator = EmbeddingSimilarityEvaluator.from_input_examples(
    examples
)
evaluator(sentence_model)
# 0.8151553052035344

 

 

 

 

3.2  미세조정 

 

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

 ↪︎  기본 임베딩 모델을 만들 때 코사인 유사도 손실을 활용해 모델을 학습시켰었는데, 이번에는 Multiple Negatives Ranking 손실을 사용해 미세조정을 시도한다. 

 ↪︎  MNR은 MRC 데이터셋과 같이 서로 관련이 있는 문장만 있는 경우 사용하기 좋은 손실함수이다. MNR 손실을 사용하면, 자동으로 하나의 배치 데이터 안에서 다른 데이터의 기사 본문을 관련없는 데이터로 사용해 모델을 학습시키기 때문에, 서로 관련이 있는 데이터만으로 학습 데이터를 구성하면 된다. 

 

1) 데이터 불러오기 

# 긍정 데이터만으로 학습 데이터 구성 
train_samples = []
for idx, row in df_train_ir.iterrows():
    train_samples.append(InputExample(
        texts=[row['question'], row['context']]
    ))
    
    
 # 중복 학습 데이터 제거 
from sentence_transformers import datasets

batch_size = 16

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

 

 

2) MNR 손실함수 불러오기 및 미세조정 

 

from sentence_transformers import losses
#[참고] sentence_model = SentenceTransformer('shangrilar/klue-roberta-base-klue-sts')

# 손실함수 불러오기 
loss = losses.MultipleNegativesRankingLoss(sentence_model)

# 미세조정 수행 
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
)

 

 

3) 평가 

 ↪︎  미세조정 전에 0.815였던 성능이 0.86으로 상승 

evaluator(sentence_model)
# 0.8600968992433692

 

 

 

4) 모델 업로드

from huggingface_hub import HfApi
api = HfApi()
repo_id = "klue-roberta-base-klue-sts-mrc"
api.create_repo(repo_id=repo_id)

api.upload_folder(
    folder_path=save_path,
    repo_id=f"본인의 아이디 입력/{repo_id}",
    repo_type="model",
)

 

 

 

 

 

4.  검색 품질을 높이는 순위 재정렬 

---

 

❍  교차인코더 미세조정

 ↪︎  교차인코더는 2개 문장을 입력받아 문장 사이의 관계를 학습하므로, 문장분류모델을 사용한다. 

 ↪︎  문장분류 모델이라, transformers 라이브러리로 모델을 직접 학습하는 방식도 가능하지만, 실습에서는 CrossEncoder와 미세조정 메서드를 사용한다. 

 

 

1) CrossEncoder 불러오기 

from sentence_transformers.cross_encoder import CrossEncoder
cross_model = CrossEncoder('klue/roberta-small', num_labels=1)

 

 ↪︎  교차인코더는 많은 계산을 해야 하므로 파라미터수가 작은 roberta-small 모델을 사용한다. 

 ↪︎  roberta-small 모델은 분류헤드가 없는 언어모델이므로, 교차인코더로 불러오려면 분류 헤드는 랜덤으로 초기화된다(성능이 낮을수밖에 없음). 

 

 

2) 초기 성능 결과

from sentence_transformers.cross_encoder.evaluation import CECorrelationEvaluator
ce_evaluator = CECorrelationEvaluator.from_input_examples(examples)
ce_evaluator(cross_model)
# 0.003316821814673943

 

 ↪︎  미세조정하지 않은 교차 인코더의 성능 결과를 보면, 문장의 관련성을 잘 계산하지 못하는 것을 확인할 수 있다. 

 

 

3) 교차인코더 학습 수행 

# 학습 데이터셋 준비 

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

 

train_batch_size = 16
num_epochs = 1
model_save_path = 'output/training_mrc'

train_dataloader = DataLoader(train_samples, shuffle=True, batch_size=train_batch_size)

cross_model.fit(
    train_dataloader=train_dataloader,
    epochs=num_epochs,
    warmup_steps=100,
    output_path=model_save_path
)


# 결과 평가 
ce_evaluator(cross_model)
# 0.8650250798639563

 

 ↪︎  학습 수행 후, 성능을 다시 확인해보면 0.865로 크게 높아진 것을 확인할 수 있다. 

 

 

 

 

5.  바이인코더와 교차 인코더로 개선된 RAG구현하기 

---

 

❍  3개 모델 학습

 ↪︎  1) 언어모델을 임베딩 모델 (사전학습된 언어모델+Pooling층) 로 변환한 기본 임베딩 모델

 ↪︎  2) 기본 임베딩 모델을 MRC 데이터셋으로 미세조정한 임베딩 모델 

 ↪︎  3) MRC 데이터셋으로 학습시킨 교차인코더 

 

❍  성능지표

 ↪︎  HitRate@10 : 질문 칼럼을 입력했을 때, 검색된 상위 10개 기사 본문에 정답이 있는 비율 

 ↪︎  관련해 evaluate_hit_rate 함수를 정의함 

 

 

5.1  기본 임베딩 모델로 검색 

 

from sentence_transformers import SentenceTransformer
base_embedding_model = SentenceTransformer('shangrilar/klue-roberta-base-klue-sts')
base_index = make_embedding_index(base_embedding_model, klue_mrc_test['context'])

evaluate_hit_rate(klue_mrc_test, base_embedding_model, base_index, 10)
# (0.88, 13.216430425643921)

 

 ↪︎  88%의 데이터에서 정답을 잘 찾았고, 평가에는 13초가 걸렸다. 총 1000개의 평가 데이터를 사용했으므로 데이터 하나당 0.013초가 소요됨 

 

 

 

5.2  미세조정한 임베딩 모델로 검색 

 

finetuned_embedding_model = SentenceTransformer('shangrilar/klue-roberta-base-klue-sts-mrc')
finetuned_index = make_embedding_index(finetuned_embedding_model, klue_mrc_test['context'])
evaluate_hit_rate(klue_mrc_test, finetuned_embedding_model, finetuned_index, 10)
# (0.946, 14.309881687164307)

 

 ↪︎  94.6%의 데이터에서 정답 기사를 정확히 가져왔다. 약간의 미세조정만으로도 7.5%정도 성능이 향상되었다. 

 

 

 

 

5.3  미세조정한 임베딩 모델과 교차 인코더 조합하기

 

❍  바이인코더 + 교차인코더

 ↪︎  임베딩 모델의 상위 N 개 결과를 받아, 교차인코더가 순위를 유사도순으로 재정렬한 후 상위 K개를 추출하면, 임베딩 모델을 통해 유사도가 높은 상위 K개를 바로 뽑았을 때보다 성능을 높일 수 있다. 

 ↪︎  교차인코더의 경우, 속도가 느리므로 전체 문서를 검색 대상으로 하지 않고 상위 N개만을 대상으로 계산하도록 범위를 좁힌다. 

 

 

import time
import numpy as np
from tqdm.auto import tqdm

# 순위 재정렬을 포함한 평가함수 
def evaluate_hit_rate_with_rerank(datasets, embedding_model, cross_model, index, bi_k=30, cross_k=10):
  start_time = time.time()
  predictions = []
  for question_idx, question in enumerate(tqdm(datasets['question'])):
    indices = find_embedding_top_k(question, embedding_model, index, bi_k)[0]
    predictions.append(rerank_top_k(cross_model, question_idx, indices, k=cross_k))
  total_prediction_count = len(predictions)
  hit_count = 0
  questions = datasets['question']
  contexts = datasets['context']
  for idx, prediction in enumerate(predictions):
    for pred in prediction:
      if contexts[pred] == contexts[idx]:
        hit_count += 1
        break
  end_time = time.time()
  return hit_count / total_prediction_count, end_time - start_time, predictions
 
 
# 결과 
hit_rate, cosumed_time, predictions = evaluate_hit_rate_with_rerank(klue_mrc_test, finetuned_embedding_model, cross_model, finetuned_index, bi_k=30, cross_k=10)
hit_rate, cosumed_time
# (0.973, 1103.055629491806)

 

 

 ↪︎  97.3% 의 데이터에서 정답을 잘 찾은 것을 볼 수 있다.