[LLM] RAG 개선하기 (3) - 바이 인코더 + 교차 인코더

 

RAG를 개선하는 최종 목적지는 결국 바이 인코더로 유사도가 높은 상위 검색 결과에 교차 인코더를 다시 한번 적용하여 검색 정확도를 최대한 끌어올리는 데에 있다. 이를 위해 이번 시간에는 교차 인코더를 구현해 보고, 이어서 바이 인코더와 교차 인코더를 조합하여 최적의 검색 기능을 만들어 보자. 그러기 위해서 앞서 작성한 포스팅들이 큰 도움이 될 것이다.

 

[LLM] 임베딩 모델로 데이터 의미 압축하기 (2) - 문장 임베딩 방식

 

교차 인코더 구현

 

교차 인코더의 A to Z를 다루는 것이 아니라 구현이라는 말이 조심스럽기는 하지만 그래도 날로 먹지는 않는다. 허깅페이스에 transformers 라이브러리로 모델을 직접 학습하는 방식도 있지만, sentence-transformers의 CrossEncoder와 미세 조정 메서드를 통해 교차 인코더를 구현해 보자.

 

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

 

먼저 roberta-small 모델에는 분류 헤드가 없다. 따라서 이걸 그대로 사용한다면 성능이 잘 나올 수가 없다.

 

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

 

이를 확인해 보기 위해 교차 인코더를 평가할 때 사용하는 CECorrelationEvaluator 클래스를 사용해 보자. 여기서의 example는 앞서 진행했던 RAG 개선하기 (2)에서 선언되고, 이를 그대로 사용하므로 해당 포스팅을 확인해 보면 된다. 결과를 확인했더니 거의 0에 가까운 성능을 보임을 알 수 있었다.

 

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

 

이제 미세 조정에 들어가야하는데, 그러기에 앞서 학습을 위한 데이터셋을 구축해 보자. 이전 포스팅의 데이터 전처리 방법과 동일한데, 질문과 내용이 대응되는 데이터에는 1을, 그렇지 못한 데이터에는 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에 가까웠던 성능이 많이 증가했음을 확인할 수 있다.

 

바이 인코더 + 교차 인코더

 

이제 교차 인코더를 다 만들어 놨으니, 바이 인코더와 교차 인코더를 모두 사용하여 개선된 RAG를 구현해 보자.

 

from datasets import load_dataset
klue_mrc_test = load_dataset('klue', 'mrc', split='validation')
klue_mrc_test = klue_mrc_test.train_test_split(test_size=1000, seed=42)['test']

 

먼저 KLUE MRC 데이터셋에서 검증 데이터셋을 가져오자. 그리고 그중 1000개만 선별하였다.

 

import faiss
def make_embedding_index(sentence_model, corpus):
	embeddings = sentence_model.encode(corpus)
	index = faiss.IndexFlatL2(embeddings.shape[1])
	index.add(embeddings)
	return index

def find_embedding_top_k(query, sentence_model, index, k):
	embedding = sentence_model.encode([query])
	distances, indices = index.search(embedding, k)
	return indices

 

이어서 두 메서드를 정의하였다. 먼저 검증 데이터셋을 문장 임베딩으로 변환하고 이를 faiss의 인덱스에 저장하는 함수, 그리고 쿼리가 주어지면 이를 임베딩 벡터로 변환 후에 앞서 정의한 index에서 찾아 유사한 k개 문서를 가져오는 함수이다.

 

def make_question_context_pairs(question_idx, indices):
  return [[klue_mrc_test['question'][question_idx], klue_mrc_test['context'][idx]] for idx in indices]

def rerank_top_k(cross_model, question_idx, indices, k):
  input_examples = make_question_context_pairs(question_idx, indices)
  relevance_scores = cross_model.predict(input_examples)
  reranked_indices = indices[np.argsort(relevance_scores)[::-1]]
  return reranked_indices

 

이제 교차 인코더의 본 기능을 수행해보자! 이를 위해  모든 질문-내용쌍을 만들어 주는 make_question_context_pairs를 선언해 준다. 이어서 이 순서쌍 데이터를 모델에 넣어 유사도 점수를 도출한 후에 유사도 점수가 높은 순으로 인덱스를 재정렬한다.

 

import time
def evaluate_hit_rate(datasets, embedding_model, index, k=10):
  start_time = time.time()
  predictions = []
  for question in datasets['question']:
    predictions.append(find_embedding_top_k(question, embedding_model, index, k)[0])
  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

 

이어지는 메서드는 hit rate를 구하기 위함이다. 얼마나 질문-내용 쌍을 잘 맞추는지에 대한 것인데, 이 함수를 실행시키면 얼마나 걸리는지와 몇 번 hit 했는지를 셀 수 있다. 먼저 질문과 유사한 k개의 답변을 뽑은 순서쌍을 predictions 리스트에 넣고, 이 중에서 정답 데이터가 포함돼 있는 경우에만 hit_count를 증가시켜 주면서 전체 예측 중에 몇 번을 맞췄는지 0 ~ 1의 값으로 반환하고, 이때 걸리는 시간을 반환한다.

 

from sentence_transformers import SentenceTransformer
base_embedding_model = SentenceTransformer('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)

 

이제 성능 측정을 위한 준비는 마쳤으니 모델을 불러와 성능을 측정해 보자. 앞서 KLUE STS 데이터로 roberta모델을 학습시켰는데, 이를 허깅페이스에 저장하지 않았었다. 이 실습을 진행하려면 앞서 만들었던 모델들을 저장해야 하는데, 만약 따라 하다가 막히는 부분이 있다면 이 부분일 것이다. 향후에 허깅페이스에 넣는 방법을 각 포스팅에 추가해 두도록 하겠다. 아무튼 불러온 모델에 테스트 데이터셋을 임베딩 변환하여 index에 넣고, 이를 바탕으로 klue_mrc_test 데이터셋에 대한 hit_rate를 알아보았다. 그 결과 0.88의 성능, 13.22 ms가 소요됨을 확인할 수 있었다.

 

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)

 

같은 방식으로 미세 조정된 모델도 불러와 본다. 이 모델 또한 앞서 학습된 모델을 허깅페이스에 저장해두었어야 불러오기가 가능할 텐데, 그러질 못했다. 아무튼 불러오는 것을 성공한다면 위와 같은 방식으로 성능 측정을 해본다. 그 결과 0.946의 성능과 14.31 ms의 시간이 소요됐음을 알 수 있다! 이를 통해 기본 모델에 비해 MRC 데이터셋으로 학습시킨 모델의 성능이 상승하였음을 확인하였다.

 

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

 

이제 진짜 마지막 단계이다. 기존의 evaluate_hit_rate() 메서드에서 두개의 인자가 더 추가되었다. bi_k는 바이 인코더로 뽑을 유사한 문서의 개수이며, cross_k는 그렇게 뽑은 문서를 바탕으로 교차 인코딩을 진행했을 때 뽑을 최종 문서의 개수를 의미한다. 첫 번째 for 문에서 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)

 

이렇게 재정렬까지 하여 모델을 평가해 보았다. 그 결과 0.973의 성능과 1,103 ms의 시간이 소요되었다! 그리고 시간은 큰 폭 증가하였음을 알 수 있다.