간단한 BERT 모델을 이용하여 GPT 챗봇 성능 높이기

 

개인 챗봇을 만들어 운영하던 도중 쉽고 재밋는 방법을 구현하게되어 공유하고자 글은 씁니다.

 

GPT의 인기가 높아지면서, 많은 기업과 개인이 언어 모델을 활용하려는 시도를 보게 됩니다. 여러 프로젝트를 진행하며, 프로젝트가 발전해가는 과정에서 종종 간과되기 쉬운, 그러나 매우 중요한 요소가 바로 '프롬프트'입니다. 프롬프트는 언어 모델을 제어하는 데 도움을 주며, 그 사용법에 따라 매우 다양한 결과를 낳을 수 있습니다.

 

그런데, 프롬프트는 프로젝트의 복잡도가 높아질수록 길어지기 마련이고, 이 때 여러 프롬프트가 서로 엉키게 되면, 오히려 간결한 프롬프트보다 효과가 떨어질 수 있습니다. 이런 문제를 해결하기 위해, 개인 챗봇에 사용되는 점점 길어지는 프롬프트를 관리하는 방법으로 '프롬프트 선택기(Prompt Selector)'를 생각해보았습니다.(인용 1) 이는 필요할 때 필요한 프롬프트만을 선택적으로 사용할 수 있게 해주는 플러그인 형태로, Bert 모델을 기반으로 구현하기로 결정했습니다." 챗봇의 간단한 아키텍처는 아래와 같습니다.

 

 

사용자 질의를 1차적으로 BERT 모델에 따라 어떤 프롬프트를 선택할지 Classify 한 후 GPT 에게 넘기는 형식

 

간단히 말하자면, 이 시스템은 사용자의 질문을 받아, BERT 분류 모델을 이용해 가장 적합한 프롬프트를 결정합니다. 이 접근 방식을 통해, 1) 응답 시간을 단축하고, 2) 답변의 정확도를 높이며, 3) 사용자가 원하는 형태의 답변을 제공할 수 있는 장점이 있습니다.

 

예를 들어, 사용자가 일반적인 질문, 회사 내부 정보에 관한 질문, 다기능을 요구하는 챗봇을 사용할 경우를 생각해 봅시다. 이런 상황에서 하나의 프롬프트로 모든 요구사항을 충족시키기는 어렵고, 프롬프트를 나누면 사용자가 수동으로 적절한 프롬프트를 선택해야 합니다. 예를 들어, 사용자가 '내일 오후 2시에 김아무개와 회의실 예약해줘'라는 질문을 했을 때, BERT 모델은 이 질문이 회의실 예약 기능을 호출해야 한다고 판단하고, 그에 맞는 프롬프트를 자동으로 선택해줍니다. 이러한 방식으로 사용자의 요구에 더 정확하고 효율적으로 대응할 수 있습니다.

 

1. 데이터 예시

# 데이터 로드
import pandas as pd
df = pd.read_csv('train.csv') 
print(df[['question', 'label']])

# 데이터와 라벨 설정
x = df.question.values
y = df.label.values

 

김아무개와 회의실 예약하고싶어	2
생성형 AI 는 뭐야?	0
우리 사내 복지는 어떤것이 있어?	1
사내 건강검진 예약방법 알려줘	1

0 = 일반질문, 1 = 사내질문, 2 = 미팅 예약

 

2. BERT Model Fine-Tuning

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification, AdamW, AutoConfig
from torch.utils.data import DataLoader, Dataset
from sklearn.model_selection import train_test_split
import numpy as np
import pandas as pd


# Load the tokenizer and model
model_name = "snunlp/KR-Medium"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Create a configuration with LoRA
config = AutoConfig.from_pretrained(model_name, 
                                    num_labels=3,
                                    lora=True,  # Enable LoRA
                                    lora_r=8,   # Rank of the LoRA matrices
                                    lora_alpha=16,  # Scaling factor for LoRA
                                    hidden_dropout_prob=0.3,  # Dropout rate for hidden layers
                                    attention_probs_dropout_prob=0.3)  # Dropout rate for attention layers

model = AutoModelForSequenceClassification.from_pretrained(model_name, config=config)

# Prepare your dataset (this is just a placeholder)
queries = x
labels = y

# Split the data into training and validation sets
train_queries, val_queries, train_labels, val_labels = train_test_split(queries, labels, test_size=0.2)

# Calculate class weights based on class distribution
class_counts = np.bincount(train_labels)
total_samples = len(train_labels)
class_weights = torch.tensor([total_samples / (2 * count) for count in class_counts], dtype=torch.float32)

# Dataset class
class QueryDataset(Dataset):
    def __init__(self, queries, labels):
        self.queries = queries
        self.labels = labels

    def __len__(self):
        return len(self.queries)

    def __getitem__(self, idx):
        query = self.queries[idx]
        label = self.labels[idx]
        inputs = tokenizer(query, padding='max_length', max_length=512, truncation=True, return_tensors="pt")
        return {
            'input_ids': inputs['input_ids'].flatten(),
            'attention_mask': inputs['attention_mask'].flatten(),
            'labels': torch.tensor(label)
        }

# Create data loaders
train_dataset = QueryDataset(train_queries, train_labels)
val_dataset = QueryDataset(val_queries, val_labels)
train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=8)

# Move model to GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# Optimizer
optimizer = AdamW(model.parameters(), lr=5e-5)

# Loss function with weighted loss
criterion = torch.nn.CrossEntropyLoss(weight=class_weights.to(device))

# Training Loop
for epoch in range(3):  # Number of epochs
    model.train()
    for batch in train_loader:
        optimizer.zero_grad()
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)
        outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
        loss = outputs.loss
        loss.backward()
        optimizer.step()

    # Validation step
    model.eval()
    val_loss = 0
    for batch in val_loader:
        with torch.no_grad():
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['labels'].to(device)
            outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
            val_loss += outputs.loss.item()
    print(f"Epoch {epoch}: Validation Loss {val_loss / len(val_loader)}")
    # Save the model after each epoch
    torch.save(model.state_dict(), f"model_epoch_{epoch}.pt")

 

특이사항에 대해 간략히 설명을 하자면, 절대적으로 적은 데이터셋을 다루기위해 weighted loss 를 계산하였으며, LORA 방식으로 fine-tuning 하였다. 

 

3. Model Inference with score

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from torch.nn.functional import softmax
# Load the Korean BERT model
model_name = "snunlp/KR-Medium"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Replace 'my_model_directory' with the actual path to your saved model directory.
model_directory = 'classification_weighted_LORA'

# Load the model from the saved directory
loaded_model = AutoModelForSequenceClassification.from_pretrained(model_directory)

# Move model to GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
loaded_model.to(device)

# Now, you can use 'loaded_model' for inference or further training.

def predict_query_with_score(model, tokenizer, query):
    # Ensure model is in evaluation mode
    model.eval()

    # Tokenize the query
    inputs = tokenizer(query, padding=True, truncation=True, max_length=512, return_tensors="pt")

    # Move inputs to the same device as the model
    input_ids = inputs['input_ids'].to(device)
    attention_mask = inputs['attention_mask'].to(device)

    # Perform inference
    with torch.no_grad():
        outputs = model(input_ids, attention_mask=attention_mask)

    # Convert logits to probabilities
    probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)

    # Get the highest probability value and the predicted class
    max_prob, predicted_class = torch.max(probabilities, dim=-1)

    return predicted_class.item(), max_prob.item()

# Example query
#query = df.question[0]
# query = "자물쇠 언락 방법"
# query = "주소 변경 방법 알려주세요"
query = "회의 예약하는법 알려줘"

# Predict the category of the query and get confidence score
predicted_category, confidence_score = predict_query_with_score(loaded_model, tokenizer, query)
print(f"The predicted category for the query '{query}' is: {predicted_category} with a confidence score of {confidence_score:.2f}")

 

 

 

 

 

인용

1. https://arxiv.org/pdf/2305.05176.pdf - FrugalGPT: How to Use Large Language Models While Reducing Cost and Improving Performance

2. https://huggingface.co/snunlp/KR-Medium - BERT Pretrained model